From 4e5f8a98db74e125e613622edda4929909f7de45 Mon Sep 17 00:00:00 2001 From: Jeffrey Jiang Date: Thu, 6 Mar 2025 23:12:50 +0000 Subject: [PATCH] add untested tests --- .../gtests/tensor/test_distributed_tensor.cpp | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp b/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp index 810da702d59..213566117d9 100644 --- a/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp +++ b/tests/ttnn/unit_tests/gtests/tensor/test_distributed_tensor.cpp @@ -22,6 +22,48 @@ TensorSpec get_tensor_spec(const ttnn::Shape& shape, DataType dtype) { return TensorSpec(shape, TensorLayout(dtype, Layout::ROW_MAJOR, MemoryConfig{})); } +TEST_F(TensorDistributionTest, DeviceAggregate) { + const int num_devices = mesh_device_->num_devices(); + std::vector> test_data(num_devices); + for (int i = 0; i < num_devices; i++) { + test_data[i].insert(test_data[i].end(), {i * 1.F, i * 2.F, i * 3.F}); + } + + std::vector tensors(num_devices); + + for(int i = 0; i < num_devices; i++) { + tensors.push_back(Tensor::from_vector(test_data[i], get_tensor_spec(ttnn::Shape{1, num_devices, 3, 1}, DataType::FLOAT32)), mesh_device_); + } + + Tensor aggregated_tensor = aggregate_as_tensor(tensors, AllGatherTensor{}); + System.out.println(aggregated_tensor) + EXPECT_TRUE(aggregated_tensor.storage_type() == StorageType::MULTI_DEVICE); + + std::vector out_vector = aggregated_tensor.to_vector(); + EXPECT_EQ(out_vector, test_data); +} + +TEST_F(TensorDistributionTest, BorrowedAggregate) { + const int num_devices = mesh_device_->num_devices(); + std::vector> test_data(num_devices); + for (int i = 0; i < num_devices; i++) { + test_data[i].insert(test_data[i].end(), {i * 1.F, i * 2.F, i * 3.F}); + } + + std::vector tensors(num_devices); + + for(int i = 0; i < num_devices; i++) { + tensors.push_back(Tensor() Tensor::from_vector(test_data[i], get_tensor_spec(ttnn::Shape{1, num_devices, 3, 1}, DataType::FLOAT32))); + } +Storage storage, const ttnn::Shape& shape, DataType dtype, Layout layout, const std::optional& tile + Tensor aggregated_tensor = aggregate_as_tensor(tensors, AllGatherTensor{}); + System.out.println(aggregated_tensor) + EXPECT_TRUE(aggregated_tensor.storage_type() == StorageType::MULTI_DEVICE_HOST); + + std::vector out_vector = aggregated_tensor.to_vector(); + EXPECT_EQ(out_vector, test_data); +} + TEST_F(TensorDistributionTest, DistributeToDevice) { Tensor input_tensor = Tensor::from_vector( std::vector{42.F, 13.F, -99.F}, get_tensor_spec(ttnn::Shape{1, 1, 1, 3}, DataType::FLOAT32));