#0: Add temporary api for automatically creating a sub-device manager…

… with a fabric sub-device. Remove create_sub_device_manager overload of MeshDevice that accepts a diffent sub-device config per device

This was added due to eth cores being different per device, but this is now handled internally

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,8 @@ 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);
+    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/device/device.cpp

@@ -3724,6 +3724,20 @@ SubDeviceManagerId Device::create_sub_device_manager(tt::stl::Span<const SubDevi
     return sub_device_manager->first;
 }
 
+std::tuple<SubDeviceManagerId, SubDeviceId> Device::create_sub_device_manager_with_fabric(tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size) {
+    auto fabric_sub_device = SubDevice(std::array{CoreRangeSet(), this->default_sub_device_manager_->sub_device(SubDeviceId{0}).cores(HalProgrammableCoreType::ACTIVE_ETH)});
+    auto new_sub_devices = std::vector<SubDevice>(sub_devices.begin(), sub_devices.end());
+    new_sub_devices.push_back(fabric_sub_device);
+    auto fabric_sub_device_id = SubDeviceId{static_cast<uint32_t>(new_sub_devices.size() - 1)};
+    auto sub_device_manager_id = this->create_sub_device_manager(new_sub_devices, local_l1_size);
+    this->sub_device_managers_[sub_device_manager_id]->set_fabric_sub_device_id(fabric_sub_device_id);
+    return {sub_device_manager_id, fabric_sub_device_id};
+}
+
+std::optional<SubDeviceId> Device::get_fabric_sub_device_id() const {
+    return this->active_sub_device_manager_->fabric_sub_device_id();
+}
+
 void Device::load_sub_device_manager(SubDeviceManagerId sub_device_manager_id) {
     TT_FATAL(!this->using_slow_dispatch(), "Using sub device managers is unsupported with slow dispatch");
     if (this->active_sub_device_manager_id_ == sub_device_manager_id) {

tt_metal/impl/device/device.hpp

@@ -389,6 +389,11 @@ class Device {
     void clear_loaded_sub_device_manager();
     void remove_sub_device_manager(SubDeviceManagerId sub_device_manager_id);
     const std::vector<SubDeviceId> &get_sub_device_ids() const;
+
+    // Temporary api until migration to actual fabric is complete
+    std::tuple<SubDeviceManagerId, SubDeviceId> create_sub_device_manager_with_fabric(tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size);
+    std::optional<SubDeviceId> get_fabric_sub_device_id() const;
+
    private:
     void initialize_default_sub_device_state(size_t l1_small_size, size_t trace_region_size, tt::stl::Span<const std::uint32_t> l1_bank_remap);
     SubDeviceManagerId get_next_sub_device_manager_id();

tt_metal/impl/sub_device/sub_device_manager.cpp

@@ -156,6 +156,16 @@ bool SubDeviceManager::has_allocations() const {
 
 DeviceAddr SubDeviceManager::local_l1_size() const { return this->local_l1_size_; }
 
+void SubDeviceManager::set_fabric_sub_device_id(SubDeviceId fabric_sub_device_id) {
+    const auto& fabric_sub_device = this->sub_device(fabric_sub_device_id);
+    TT_FATAL(
+        fabric_sub_device.cores(HalProgrammableCoreType::TENSIX).num_cores() == 0,
+        "Fabric sub device must not have Tensix cores");
+    this->fabric_sub_device_id_ = fabric_sub_device_id;
+}
+
+std::optional<SubDeviceId> SubDeviceManager::fabric_sub_device_id() const { return this->fabric_sub_device_id_; }
+
 uint8_t SubDeviceManager::get_sub_device_index(SubDeviceId sub_device_id) const {
     auto sub_device_index = sub_device_id.to_index();
     TT_FATAL(

tt_metal/impl/sub_device/sub_device_manager.hpp

@@ -68,6 +68,10 @@ class SubDeviceManager {
     bool has_allocations() const;
     DeviceAddr local_l1_size() const;
 
+    // Temporary until migration to actual fabric is complete
+    void set_fabric_sub_device_id(SubDeviceId sub_device_id);
+    std::optional<SubDeviceId> fabric_sub_device_id() const;
+
 private:
     void validate_sub_devices() const;
     uint8_t get_sub_device_index(SubDeviceId sub_device_id) const;
@@ -97,6 +101,9 @@ class SubDeviceManager {
     std::unordered_map<uint32_t, std::shared_ptr<TraceBuffer>> trace_buffer_pool_;
 
     std::vector<LaunchMessageRingBufferState> worker_launch_message_buffer_state_;
+
+    // Temporary until migration to actual fabric is complete
+    std::optional<SubDeviceId> fabric_sub_device_id_ = std::nullopt;
 };
 
 }  // namespace detail

ttnn/cpp/ttnn/distributed/distributed_pybind.cpp

@@ -197,23 +197,24 @@ void py_module(py::module& module) {
                     MeshSubDeviceManagerId: The ID of the created sub-device manager.
             )doc")
         .def(
-            "create_sub_device_manager",
-            [](MeshDevice& self,
-               const std::vector<std::vector<SubDevice>>& mesh_sub_devices,
-               DeviceAddr local_l1_size) { return self.create_sub_device_manager(mesh_sub_devices, local_l1_size); },
+            "create_sub_device_manager_with_fabric",
+            [](MeshDevice& self, const std::vector<SubDevice>& sub_devices, DeviceAddr local_l1_size) {
+                return self.create_sub_device_manager(sub_devices, local_l1_size);
+            },
             py::arg("sub_devices"),
             py::arg("local_l1_size"),
             R"doc(
-                Creates a sub-device manager for the given mesh device.
+                Creates a sub-device manager for the given mesh device. This will automatically create a sub-device of ethernet cores for use with fabric.
+                Note that this is a temporary API until migration to actual fabric is complete.
 
                 Args:
-                    mesh_sub_devices (List[List[ttnn.SubDevice]]): The sub-devices to include in the sub-device manager.
-                    Each element of the outer list will be used to configure the corresponding device in the MeshDevice.
-                    This means that the individual devices in the MeshDevice may have different configurations.
+                    sub_devices (List[ttnn.SubDevice]): The sub-devices to include in the sub-device manager. No ethernet cores should be included in this list.
+                    This configuration will be used for each device in the MeshDevice.
                     local_l1_size (int): The size of the local allocators of each sub-device. The global allocator will be shrunk by this amount.
 
                 Returns:
                     MeshSubDeviceManagerId: The ID of the created sub-device manager.
+                    SubDeviceId: The ID of the sub-device that will be used for fabric.
             )doc")
         .def(
             "load_sub_device_manager",