#15297: Allow MeshDevice to be initialized for chips without eth coor…

…dinates
tenstorrent · Nov 28, 2024 · ec4bd12 · ec4bd12
1 parent 502d364
commit ec4bd12
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 2 deletions.
diff --git a/tests/scripts/run_cpp_unit_tests.sh b/tests/scripts/run_cpp_unit_tests.sh
@@ -19,6 +19,7 @@ rm -rf $kernel_path
 ./build/test/tt_metal/unit_tests_eth
 ./build/test/tt_metal/unit_tests_llk
 ./build/test/tt_metal/unit_tests_stl
+./build/test/tt_metal/distributed/distributed_unit_tests --gtest_filter=MeshDeviceSuite.*
 
 if [[ ! -z "$TT_METAL_SLOW_DISPATCH_MODE" ]]; then
     env python tests/scripts/run_tt_metal.py --dispatch-mode slow

diff --git a/tests/tt_metal/distributed/test_distributed.cpp b/tests/tt_metal/distributed/test_distributed.cpp
@@ -42,4 +42,16 @@ TEST_F(MeshDevice_T3000, SimpleMeshDeviceTest) {
     EXPECT_EQ(mesh_device_->num_cols(), 4);
 }
 
+TEST(MeshDeviceSuite, Test1x1SystemMeshInitialize) {
+    auto& sys = tt::tt_metal::distributed::SystemMesh::instance();
+
+    auto config = tt::tt_metal::distributed::MeshDeviceConfig
+        ({1, 1}, std::pair<size_t, size_t>(0, 0), {}, MeshType::RowMajor);
+
+    EXPECT_NO_THROW({
+        auto mesh = tt::tt_metal::distributed::MeshDevice::create(config, DEFAULT_L1_SMALL_SIZE, DEFAULT_TRACE_REGION_SIZE, 1,  tt::tt_metal::DispatchCoreType::WORKER);
+        mesh->close_devices();
+    });
+}
+
 }  // namespace tt::tt_metal::distributed::test
diff --git a/tt_metal/distributed/mesh_device.cpp b/tt_metal/distributed/mesh_device.cpp
@@ -100,8 +100,20 @@ SystemMesh& SystemMesh::instance() {
 }
 void SystemMesh::initialize() {
     this->physical_device_id_to_coordinate = tt::Cluster::instance().get_user_chip_ethernet_coordinates();
-    for (const auto& [chip_id, physical_coordinate] : this->physical_device_id_to_coordinate) {
-        this->physical_coordinate_to_device_id.emplace(physical_coordinate, chip_id);
+    if (this->physical_device_id_to_coordinate.empty()) {
+        // Only WH has ethernet coordinates. Fabric will assign chip ids for BH
+        auto arch = tt::Cluster::instance().arch();
+        TT_FATAL(arch == ARCH::GRAYSKULL or arch == ARCH::BLACKHOLE, "Expected Wormhole chips to have ethernet coordinates assigned by cluster descriptor");
+        const int num_detected_devices = tt::Cluster::instance().number_of_devices();
+        for (auto chip_id = 0; chip_id < num_detected_devices; chip_id++) {
+            PhysicalCoordinate coord{0, chip_id, 0, 0, 0};
+            this->physical_device_id_to_coordinate.emplace(chip_id, coord);
+            this->physical_coordinate_to_device_id.emplace(coord, chip_id);
+        }
+    } else {
+        for (const auto& [chip_id, physical_coordinate] : this->physical_device_id_to_coordinate) {
+            this->physical_coordinate_to_device_id.emplace(physical_coordinate, chip_id);
+        }
     }
 
     // Initialize the system mesh shape and translation map