Sub-Device fabric api updates, and make global sem/cb thread safe in metal

tests/tt_metal/tt_metal/api/test_global_semaphores.cpp

@@ -86,6 +86,7 @@ TEST_F(DispatchFixture, ResetGlobalSemaphores) {
     for (auto device : devices_) {
         {
             uint32_t initial_value = 1;
+            uint32_t reset_value = 2;
             std::vector<uint32_t> overwrite_value = {2};
             auto global_semaphore = tt::tt_metal::CreateGlobalSemaphore(device, cores, initial_value);
             auto address = global_semaphore->address();
@@ -104,14 +105,14 @@ TEST_F(DispatchFixture, ResetGlobalSemaphores) {
 
                 EXPECT_EQ(sem_vals[0], overwrite_value[0]);
             }
-            global_semaphore->reset_semaphore_value();
+            global_semaphore->reset_semaphore_value(reset_value);
             Synchronize(device);
             for (const auto& core : cores_vec) {
                 auto sem_vals = tt::llrt::read_hex_vec_from_core(
                     device->id(), device->worker_core_from_logical_core(core), address, sizeof(uint32_t));
                 tt::llrt::write_hex_vec_to_core(
                     device->id(), device->worker_core_from_logical_core(core), overwrite_value, address);
-                EXPECT_EQ(sem_vals[0], initial_value);
+                EXPECT_EQ(sem_vals[0], reset_value);
             }
         }
     }

tests/ttnn/unit_tests/test_global_semaphore.py

@@ -29,7 +29,7 @@ def run_global_semaphore(device):
 
     assert ttnn.get_global_semaphore_address(global_sem0) != ttnn.get_global_semaphore_address(global_sem1)
 
-    ttnn.reset_global_semaphore_value(global_sem0)
+    ttnn.reset_global_semaphore_value(global_sem0, 3)
 
 
 @pytest.mark.parametrize("enable_async_mode", (False, True), indirect=True)

tests/ttnn/unit_tests/test_sub_device.py

@@ -7,7 +7,7 @@
 import ttnn
 
 
-def run_sub_devices(device, replicate_sub_devices=False):
+def run_sub_devices(device, create_fabric_sub_device=False):
     tensix_cores0 = ttnn.CoreRangeSet(
         {
             ttnn.CoreRange(
@@ -28,12 +28,12 @@ def run_sub_devices(device, replicate_sub_devices=False):
     sub_device_2 = ttnn.SubDevice([tensix_cores1])
     sub_devices_1 = [sub_device_1, sub_device_2]
     sub_devices_2 = [sub_device_2]
-    if replicate_sub_devices:
-        num_devices = 1 if isinstance(device, ttnn.Device) else device.get_num_devices()
-        sub_devices_1 = [sub_devices_1] * num_devices
-        sub_devices_2 = [sub_devices_2] * num_devices
-    sub_device_manager1 = device.create_sub_device_manager(sub_devices_1, 3200)
-    sub_device_manager2 = device.create_sub_device_manager(sub_devices_2, 3200)
+    if create_fabric_sub_device:
+        sub_device_manager1 = device.create_sub_device_manager_with_fabric(sub_devices_1, 3200)
+        sub_device_manager2 = device.create_sub_device_manager_with_fabric(sub_devices_2, 3200)
+    else:
+        sub_device_manager1 = device.create_sub_device_manager(sub_devices_1, 3200)
+        sub_device_manager2 = device.create_sub_device_manager(sub_devices_2, 3200)
     device.load_sub_device_manager(sub_device_manager1)
     ttnn.synchronize_devices(device, sub_device_ids=[ttnn.SubDeviceId(1)])
     ttnn.synchronize_devices(device, sub_device_ids=[ttnn.SubDeviceId(0), ttnn.SubDeviceId(1)])
@@ -45,7 +45,7 @@ def run_sub_devices(device, replicate_sub_devices=False):
     device.remove_sub_device_manager(sub_device_manager2)
 
 
-def run_sub_devices_program(device, replicate_sub_devices=False):
+def run_sub_devices_program(device, create_fabric_sub_device=False):
     is_mesh_device = isinstance(device, ttnn.MeshDevice)
     if is_mesh_device:
         inputs_mesh_mapper = ttnn.ShardTensorToMesh(device, dim=0)
@@ -74,10 +74,10 @@ def run_sub_devices_program(device, replicate_sub_devices=False):
     sub_device_1 = ttnn.SubDevice([tensix_cores0])
     sub_device_2 = ttnn.SubDevice([tensix_cores1])
     sub_devices = [sub_device_1, sub_device_2]
-    if replicate_sub_devices:
-        num_devices = 1 if isinstance(device, ttnn.Device) else device.get_num_devices()
-        sub_devices = [sub_devices] * num_devices
-    sub_device_manager = device.create_sub_device_manager(sub_devices, 3200)
+    if create_fabric_sub_device:
+        sub_device_manager = device.create_sub_device_manager_with_fabric(sub_devices, 3200)
+    else:
+        sub_device_manager = device.create_sub_device_manager(sub_devices, 3200)
     device.load_sub_device_manager(sub_device_manager)
 
     x = torch.randn(num_devices, 1, 64, 64, dtype=torch.bfloat16)
@@ -140,9 +140,9 @@ def test_sub_devices(device, enable_async_mode):
 
 
 @pytest.mark.parametrize("enable_async_mode", (False, True), indirect=True)
-@pytest.mark.parametrize("replicate_sub_devices", (False, True))
-def test_sub_devices_mesh(mesh_device, replicate_sub_devices, enable_async_mode):
-    run_sub_devices(mesh_device, replicate_sub_devices)
+@pytest.mark.parametrize("create_fabric_sub_device", (False, True))
+def test_sub_devices_mesh(mesh_device, create_fabric_sub_device, enable_async_mode):
+    run_sub_devices(mesh_device, create_fabric_sub_device)
 
 
 @pytest.mark.parametrize("enable_async_mode", (False, True), indirect=True)
@@ -151,6 +151,6 @@ def test_sub_device_program(device, enable_async_mode):
 
 
 @pytest.mark.parametrize("enable_async_mode", (False, True), indirect=True)
-@pytest.mark.parametrize("replicate_sub_devices", (False, True))
-def test_sub_device_program_mesh(mesh_device, replicate_sub_devices, enable_async_mode):
-    run_sub_devices_program(mesh_device, replicate_sub_devices)
+@pytest.mark.parametrize("create_fabric_sub_device", (False, True))
+def test_sub_device_program_mesh(mesh_device, create_fabric_sub_device, enable_async_mode):
+    run_sub_devices_program(mesh_device, create_fabric_sub_device)

tt_metal/distributed/mesh_device.cpp

@@ -490,21 +490,21 @@ MeshSubDeviceManagerId MeshDevice::create_sub_device_manager(tt::stl::Span<const
     return mesh_sub_device_manager_id;
 }
 
-MeshSubDeviceManagerId MeshDevice::create_sub_device_manager(const std::vector<std::vector<SubDevice>>& mesh_sub_devices, DeviceAddr local_l1_size) {
+std::tuple<MeshSubDeviceManagerId, SubDeviceId> MeshDevice::create_sub_device_manager_with_fabric(tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size) {
     MeshSubDeviceManagerId mesh_sub_device_manager_id(*this);
-    TT_FATAL(mesh_sub_devices.size() == this->num_devices(), "Number of devices does not match number of sub-device configurations");
+    SubDeviceId fabric_sub_device_id;
     for (uint32_t i = 0; i < this->num_devices(); i++) {
         auto* device = this->devices[i];
         auto& sub_device_manager_id = mesh_sub_device_manager_id.sub_device_manager_ids[i];
-        tt::stl::Span<const SubDevice> sub_devices(mesh_sub_devices[i]);
-        device->push_work([device, sub_devices, local_l1_size, &sub_device_manager_id]() {
-            sub_device_manager_id = device->create_sub_device_manager(sub_devices, local_l1_size);
+        // All fabric sub-device ids will be the same, since all managers are created with the same sub_devices input
+        device->push_work([device, sub_devices, local_l1_size, &sub_device_manager_id, &fabric_sub_device_id]() {
+            std::tie(sub_device_manager_id, fabric_sub_device_id) = device->create_sub_device_manager_with_fabric(sub_devices, local_l1_size);
         });
     }
     for (auto* device : this->devices) {
         device->synchronize();
     }
-    return mesh_sub_device_manager_id;
+    return {mesh_sub_device_manager_id, fabric_sub_device_id};
 }
 
 void MeshDevice::load_sub_device_manager(MeshSubDeviceManagerId mesh_sub_device_manager_id) {

tt_metal/distributed/mesh_device.hpp

@@ -137,8 +137,9 @@ class MeshDevice : public std::enable_shared_from_this<MeshDevice> {
 
     MeshSubDeviceManagerId create_sub_device_manager(
         tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size);
-    MeshSubDeviceManagerId create_sub_device_manager(
-        const std::vector<std::vector<SubDevice>>& mesh_sub_devices, DeviceAddr local_l1_size);
+    // TODO #15944: Temporary api until migration to actual fabric is complete
+    std::tuple<MeshSubDeviceManagerId, SubDeviceId> create_sub_device_manager_with_fabric(
+        tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size);
     void load_sub_device_manager(MeshSubDeviceManagerId mesh_sub_device_manager_id);
     void clear_loaded_sub_device_manager();
     void remove_sub_device_manager(MeshSubDeviceManagerId mesh_sub_device_manager_id);

tt_metal/impl/buffers/global_circular_buffer.cpp

@@ -28,7 +28,8 @@ GlobalCircularBuffer::GlobalCircularBuffer(
     const std::unordered_map<CoreCoord, CoreRangeSet>& sender_receiver_core_mapping,
     uint32_t size,
     BufferType buffer_type,
-    tt::stl::Span<const SubDeviceId> sub_device_ids) :
+    tt::stl::Span<const SubDeviceId> sub_device_ids,
+    Private) :
     device_(device), sender_receiver_core_mapping_(sender_receiver_core_mapping), size_(size) {
     TT_FATAL(this->device_ != nullptr, "Device cannot be null");
     uint32_t num_sender_cores = sender_receiver_core_mapping.size();
@@ -86,58 +87,65 @@ void GlobalCircularBuffer::setup_cb_buffers(
         shard_parameters,
         std::nullopt);
 
-    const auto& core_to_core_id = this->cb_config_buffer_->get_buffer_page_mapping()->core_to_core_id_;
-
-    std::vector<uint32_t> cb_config_host_buffer(cb_config_size / sizeof(uint32_t), 0);
-    uint32_t buffer_address = this->cb_buffer_->address();
-    uint32_t noc_xy_address = this->cb_config_buffer_->address() + num_config_elements * sizeof(uint32_t);
-    uint32_t pages_sent_address = align(noc_xy_address + num_noc_xy_words * sizeof(uint32_t), l1_alignment);
-
-    for (const auto& [sender_core, receiver_cores] : this->sender_receiver_core_mapping_) {
-        const auto& receiver_cores_vec = corerange_to_cores(receiver_cores);
-        uint32_t sender_idx = core_to_core_id.at(sender_core) * cb_config_page_size / sizeof(uint32_t);
-        uint32_t num_receivers = receiver_cores.num_cores();
-        uint32_t pages_acked_address = pages_sent_address + num_receivers * l1_alignment;
-        cb_config_host_buffer[sender_idx++] = 1;
-        cb_config_host_buffer[sender_idx++] = receiver_cores.num_cores();
-        cb_config_host_buffer[sender_idx++] = buffer_address;
-        cb_config_host_buffer[sender_idx++] = this->size_;
-        cb_config_host_buffer[sender_idx++] = buffer_address;
-        cb_config_host_buffer[sender_idx++] = noc_xy_address;
-        cb_config_host_buffer[sender_idx++] = pages_sent_address;
-
-        auto sender_physical_coord = this->device_->worker_core_from_logical_core(sender_core);
-        for (uint32_t i = 0; i < receiver_cores_vec.size(); i++) {
-            auto receiver_physical_coord = this->device_->worker_core_from_logical_core(receiver_cores_vec[i]);
-            cb_config_host_buffer[sender_idx++] = receiver_physical_coord.x;
-            cb_config_host_buffer[sender_idx++] = receiver_physical_coord.y;
-
-            uint32_t receiver_idx = core_to_core_id.at(receiver_cores_vec[i]) * cb_config_page_size / sizeof(uint32_t);
-            cb_config_host_buffer[receiver_idx++] = 0;
-            cb_config_host_buffer[receiver_idx++] = num_receivers;
-            cb_config_host_buffer[receiver_idx++] = buffer_address;
-            cb_config_host_buffer[receiver_idx++] = this->size_;
-            cb_config_host_buffer[receiver_idx++] = buffer_address;
-            cb_config_host_buffer[receiver_idx++] = noc_xy_address;
-            cb_config_host_buffer[receiver_idx++] = pages_sent_address + 2 * i * l1_alignment;
-            cb_config_host_buffer[receiver_idx++] = sender_physical_coord.x;
-            cb_config_host_buffer[receiver_idx++] = sender_physical_coord.y;
-        }
-    }
-
     // Write the config buffer to the device
     // Only block for the slow dispatch case
-    if (this->device_->using_slow_dispatch()) {
-        detail::WriteToBuffer(*this->cb_config_buffer_, cb_config_host_buffer);
-        tt::Cluster::instance().l1_barrier(this->device_->id());
-    } else {
-        EnqueueWriteBuffer(
-            this->device_->command_queue(),
-            this->cb_config_buffer_,
-            cb_config_host_buffer.data(),
-            false,
-            sub_device_ids);
-    }
+    auto* device = this->device_;
+    device->push_work([device,
+                       cb_config_size,
+                       cb_config_page_size,
+                       num_noc_xy_words,
+                       l1_alignment,
+                       buffer_address = this->cb_buffer_->address(),
+                       cb_config_buffer = this->cb_config_buffer_,
+                       size = this->size_,
+                       sender_receiver_core_mapping = this->sender_receiver_core_mapping_,
+                       sub_device_ids = std::vector<SubDeviceId>(sub_device_ids.begin(), sub_device_ids.end())] {
+        auto config_buffer_address = cb_config_buffer->address();
+        const auto& core_to_core_id = cb_config_buffer->get_buffer_page_mapping()->core_to_core_id_;
+        std::vector<uint32_t> cb_config_host_buffer(cb_config_size / sizeof(uint32_t), 0);
+        uint32_t noc_xy_address = config_buffer_address + num_config_elements * sizeof(uint32_t);
+        uint32_t pages_sent_address = align(noc_xy_address + num_noc_xy_words * sizeof(uint32_t), l1_alignment);
+
+        for (const auto& [sender_core, receiver_cores] : sender_receiver_core_mapping) {
+            const auto& receiver_cores_vec = corerange_to_cores(receiver_cores);
+            uint32_t sender_idx = core_to_core_id.at(sender_core) * cb_config_page_size / sizeof(uint32_t);
+            uint32_t num_receivers = receiver_cores.num_cores();
+            uint32_t pages_acked_address = pages_sent_address + num_receivers * l1_alignment;
+            cb_config_host_buffer[sender_idx++] = 1;
+            cb_config_host_buffer[sender_idx++] = receiver_cores.num_cores();
+            cb_config_host_buffer[sender_idx++] = buffer_address;
+            cb_config_host_buffer[sender_idx++] = size;
+            cb_config_host_buffer[sender_idx++] = buffer_address;
+            cb_config_host_buffer[sender_idx++] = noc_xy_address;
+            cb_config_host_buffer[sender_idx++] = pages_sent_address;
+
+            auto sender_physical_coord = device->worker_core_from_logical_core(sender_core);
+            for (uint32_t i = 0; i < receiver_cores_vec.size(); i++) {
+                auto receiver_physical_coord = device->worker_core_from_logical_core(receiver_cores_vec[i]);
+                cb_config_host_buffer[sender_idx++] = receiver_physical_coord.x;
+                cb_config_host_buffer[sender_idx++] = receiver_physical_coord.y;
+
+                uint32_t receiver_idx =
+                    core_to_core_id.at(receiver_cores_vec[i]) * cb_config_page_size / sizeof(uint32_t);
+                cb_config_host_buffer[receiver_idx++] = 0;
+                cb_config_host_buffer[receiver_idx++] = num_receivers;
+                cb_config_host_buffer[receiver_idx++] = buffer_address;
+                cb_config_host_buffer[receiver_idx++] = size;
+                cb_config_host_buffer[receiver_idx++] = buffer_address;
+                cb_config_host_buffer[receiver_idx++] = noc_xy_address;
+                cb_config_host_buffer[receiver_idx++] = pages_sent_address + 2 * i * l1_alignment;
+                cb_config_host_buffer[receiver_idx++] = sender_physical_coord.x;
+                cb_config_host_buffer[receiver_idx++] = sender_physical_coord.y;
+            }
+        }
+        if (device->using_slow_dispatch()) {
+            detail::WriteToBuffer(*cb_config_buffer, cb_config_host_buffer);
+            tt::Cluster::instance().l1_barrier(device->id());
+        } else {
+            EnqueueWriteBuffer(
+                device->command_queue(), cb_config_buffer, cb_config_host_buffer.data(), false, sub_device_ids);
+        }
+    });
 }
 
 std::shared_ptr<GlobalCircularBuffer> GlobalCircularBuffer::create(
@@ -147,7 +155,7 @@ std::shared_ptr<GlobalCircularBuffer> GlobalCircularBuffer::create(
     BufferType buffer_type,
     tt::stl::Span<const SubDeviceId> sub_device_ids) {
     return std::make_shared<GlobalCircularBuffer>(
-        device, sender_receiver_core_mapping, size, buffer_type, sub_device_ids);
+        device, sender_receiver_core_mapping, size, buffer_type, sub_device_ids, Private());
 }
 
 const Buffer& GlobalCircularBuffer::cb_buffer() const { return *this->cb_buffer_; }

tt_metal/impl/buffers/global_circular_buffer.hpp

@@ -26,27 +26,24 @@ namespace v1 {
 namespace experimental {
 
 class GlobalCircularBuffer {
-public:
-    GlobalCircularBuffer(
-        Device* device,
-        const std::unordered_map<CoreCoord, CoreRangeSet>& sender_receiver_core_mapping,
-        uint32_t size,
-        BufferType buffer_type,
-        tt::stl::Span<const SubDeviceId> sub_device_ids);
-
-    GlobalCircularBuffer(const GlobalCircularBuffer&) = default;
-    GlobalCircularBuffer& operator=(const GlobalCircularBuffer&) = default;
-
-    GlobalCircularBuffer(GlobalCircularBuffer&&) noexcept = default;
-    GlobalCircularBuffer& operator=(GlobalCircularBuffer&&) noexcept = default;
+    struct Private {
+        explicit Private() = default;
+    };
 
+public:
     static std::shared_ptr<GlobalCircularBuffer> create(
         Device* device,
         const std::unordered_map<CoreCoord, CoreRangeSet>& sender_receiver_core_mapping,
         uint32_t size,
         BufferType buffer_type = BufferType::L1,
         tt::stl::Span<const SubDeviceId> sub_device_ids = {});
 
+    GlobalCircularBuffer(const GlobalCircularBuffer&) = delete;
+    GlobalCircularBuffer& operator=(const GlobalCircularBuffer&) = delete;
+
+    GlobalCircularBuffer(GlobalCircularBuffer&&) noexcept = delete;
+    GlobalCircularBuffer& operator=(GlobalCircularBuffer&&) noexcept = delete;
+
     const Buffer& cb_buffer() const;
 
     const CoreRangeSet& sender_cores() const;
@@ -59,6 +56,16 @@ class GlobalCircularBuffer {
     static constexpr auto attribute_names = std::forward_as_tuple("sender_receiver_core_mapping", "size");
     const auto attribute_values() const { return std::make_tuple(this->sender_receiver_core_mapping_, this->size_); }
 
+    // "Private" constructor to prevent direct instantiation
+    // Use GlobalCircularBuffer::create instead
+    GlobalCircularBuffer(
+        Device* device,
+        const std::unordered_map<CoreCoord, CoreRangeSet>& sender_receiver_core_mapping,
+        uint32_t size,
+        BufferType buffer_type,
+        tt::stl::Span<const SubDeviceId> sub_device_ids,
+        Private);
+
 private:
     void setup_cb_buffers(
         BufferType buffer_type, uint32_t max_num_receivers_per_sender, tt::stl::Span<const SubDeviceId> sub_device_ids);