#0: Add borrowed storage support for the aggregate_as_tensor function

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<std::vector<float>> 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<Tensor> 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<float> out_vector = aggregated_tensor.to_vector<float>();
+    EXPECT_EQ(out_vector, test_data);
+}
+
+TEST_F(TensorDistributionTest, BorrowedAggregate) {
+    const int num_devices = mesh_device_->num_devices();
+    std::vector<std::vector<float>> 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<Tensor> 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>& 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<float> out_vector = aggregated_tensor.to_vector<float>();
+    EXPECT_EQ(out_vector, test_data);
+}
+
 TEST_F(TensorDistributionTest, DistributeToDevice) {
     Tensor input_tensor = Tensor::from_vector(
         std::vector<float>{42.F, 13.F, -99.F}, get_tensor_spec(ttnn::Shape{1, 1, 1, 3}, DataType::FLOAT32));

ttnn/cpp/ttnn/distributed/api.cpp

@@ -7,14 +7,17 @@
 #include <memory>
 
 #include <tt-metalium/overloaded.hpp>
+#include "tt-metalium/assert.hpp"
 #include "tt-metalium/mesh_coord.hpp"
 #include "ttnn/tensor/tensor.hpp"
+#include "ttnn/tensor/host_buffer/functions.hpp"
 #include "ttnn/tensor/tensor_utils.hpp"
 #include "ttnn/distributed/distributed_tensor_config.hpp"
 #include <tt-metalium/mesh_device.hpp>
 #include <tt-metalium/system_mesh.hpp>
 #include "ttnn/distributed/distributed_tensor_config.hpp"
 
+
 using namespace tt::tt_metal;
 
 namespace ttnn::distributed {
@@ -93,7 +96,25 @@ Tensor aggregate_as_tensor(
         }
         auto storage = MultiDeviceHostStorage{config, std::move(host_owned_buffers), specs};
         return Tensor(std::move(storage), reference_shard.get_tensor_spec());
+} else if (storage_type == StorageType::BORROWED) {
+        std::vector<ttnn::TensorSpec> specs;
+        std::vector<OwnedBuffer> host_owned_buffers;
+        for (const auto& shard : tensor_shards) {
+            auto buffer = std::get<BorrowedStorage>(shard.get_storage()).buffer;
+            specs.push_back(shard.get_tensor_spec());
+
+            auto visitor = tt::stl::overloaded{[&shard, &host_owned_buffers](const auto& buffer) -> OwnedBuffer {
+                using BorrowedBufferType = std::vector<typename std::decay_t<decltype(buffer)>::value_type>;
+
+                return owned_buffer::create(BorrowedBufferType(buffer.begin(), buffer.end()));
+            }};
+
+            host_owned_buffers.push_back(std::visit(visitor, buffer));
+        }
+        auto storage = MultiDeviceHostStorage{config, std::move(host_owned_buffers), specs};
+        return Tensor(std::move(storage), reference_shard.get_tensor_spec());
     } else {
+        TT_FATAL(storage_type == StorageType::DEVICE, "Unexpected storage type {}", storage_type);
         std::vector<int> ordered_device_ids;
         std::unordered_map<int, ttnn::TensorSpec> specs;
         std::unordered_map<int, std::shared_ptr<Buffer>> device_buffers;