Blackhole DRAM harvesting

tenstorrent · Nov 29, 2024 · c078977 · c078977
1 parent 4a0c5c7
commit c078977
Show file tree

Hide file tree

Showing 6 changed files with 141 additions and 13 deletions.
diff --git a/device/api/umd/device/blackhole_coordinate_manager.h b/device/api/umd/device/blackhole_coordinate_manager.h
@@ -32,4 +32,8 @@ class BlackholeCoordinateManager : public CoordinateManager {
     void fill_tensix_logical_to_translated() override;
     void fill_eth_logical_to_translated() override;
     void fill_pcie_logical_to_translated() override;
+    void fill_dram_logical_to_translated() override;
+
+private:
+    void map_column_of_dram_banks(const size_t start_bank, const size_t end_bank, const size_t x_coord);
 };
diff --git a/device/api/umd/device/blackhole_implementation.h b/device/api/umd/device/blackhole_implementation.h
@@ -180,6 +180,9 @@ static const size_t eth_translated_coordinate_start_y = 25;
 static const size_t pcie_translated_coordinate_start_x = 19;
 static const size_t pcie_translated_coordinate_start_y = 24;
 
+static const size_t dram_translated_coordinate_start_x = 17;
+static const size_t dram_translated_coordinate_start_y = 12;
+
 }  // namespace blackhole
 
 class blackhole_implementation : public architecture_implementation {

diff --git a/device/api/umd/device/coordinate_manager.h b/device/api/umd/device/coordinate_manager.h
@@ -50,6 +50,8 @@ class CoordinateManager {
 
     static size_t get_num_harvested(const size_t harvesting_mask);
 
+    static std::vector<size_t> get_harvested_indices(const size_t harvesting_mask);
+
     CoordinateManager(CoordinateManager& other) = default;
 
     tt::umd::CoreCoord to(const tt::umd::CoreCoord core_coord, const CoordSystem coord_system);
@@ -62,6 +64,9 @@ class CoordinateManager {
     tt::umd::CoreCoord to_virtual(const tt::umd::CoreCoord core_coord);
     tt::umd::CoreCoord to_translated(const tt::umd::CoreCoord core_coord);
 
+    static void assert_create_coordinate_manager(
+        const tt::ARCH arch, const size_t tensix_harvesting_mask, const size_t dram_harvesting_mask);
+
 protected:
     virtual void translate_tensix_coords();
     virtual void translate_dram_coords();

diff --git a/device/blackhole/blackhole_coordinate_manager.cpp b/device/blackhole/blackhole_coordinate_manager.cpp
@@ -112,6 +112,8 @@ void BlackholeCoordinateManager::translate_dram_coords() {
             logical_x++;
         }
     }
+
+    fill_dram_logical_to_translated();
 }
 
 void BlackholeCoordinateManager::fill_eth_logical_to_translated() {
@@ -137,3 +139,49 @@ void BlackholeCoordinateManager::fill_pcie_logical_to_translated() {
         blackhole::pcie_translated_coordinate_start_x, blackhole::pcie_translated_coordinate_start_y}] =
         CoreCoord(0, 0, CoreType::PCIE, CoordSystem::LOGICAL);
 }
+
+void BlackholeCoordinateManager::map_column_of_dram_banks(
+    const size_t start_bank, const size_t end_bank, const size_t x_coord) {
+    size_t translated_y = blackhole::dram_translated_coordinate_start_y;
+    for (size_t bank = start_bank; bank < end_bank; bank++) {
+        for (size_t port = 0; port < blackhole::NUM_NOC_PORTS_PER_DRAM_BANK; port++) {
+            dram_logical_to_translated[{bank, port}] =
+                CoreCoord(x_coord, translated_y, CoreType::DRAM, CoordSystem::TRANSLATED);
+            dram_translated_to_logical[{x_coord, translated_y}] =
+                CoreCoord(bank, port, CoreType::DRAM, CoordSystem::LOGICAL);
+            translated_y++;
+        }
+    }
+}
+
+void BlackholeCoordinateManager::fill_dram_logical_to_translated() {
+    const std::vector<size_t> harvested_banks = CoordinateManager::get_harvested_indices(dram_harvesting_mask);
+
+    if (harvested_banks.empty()) {
+        map_column_of_dram_banks(0, blackhole::NUM_DRAM_BANKS / 2, blackhole::dram_translated_coordinate_start_x);
+        map_column_of_dram_banks(
+            blackhole::NUM_DRAM_BANKS / 2,
+            blackhole::NUM_DRAM_BANKS,
+            blackhole::dram_translated_coordinate_start_x + 1);
+        return;
+    }
+
+    const size_t harvested_bank = harvested_banks[0];
+
+    if (harvested_bank < blackhole::NUM_DRAM_BANKS / 2) {
+        const size_t mirror_east_bank = harvested_bank + blackhole::NUM_DRAM_BANKS / 2;
+        map_column_of_dram_banks(
+            0, blackhole::NUM_DRAM_BANKS / 2 - 1, blackhole::dram_translated_coordinate_start_x + 1);
+        map_column_of_dram_banks(
+            blackhole::NUM_DRAM_BANKS / 2 - 1,
+            blackhole::NUM_DRAM_BANKS - 1,
+            blackhole::dram_translated_coordinate_start_x);
+    } else {
+        const size_t mirror_west_bank = harvested_bank - blackhole::NUM_DRAM_BANKS / 2;
+        map_column_of_dram_banks(0, blackhole::NUM_DRAM_BANKS / 2, blackhole::dram_translated_coordinate_start_x);
+        map_column_of_dram_banks(
+            blackhole::NUM_DRAM_BANKS / 2,
+            blackhole::NUM_DRAM_BANKS - 1,
+            blackhole::dram_translated_coordinate_start_x + 1);
+    }
+}
diff --git a/device/coordinate_manager.cpp b/device/coordinate_manager.cpp
@@ -395,11 +395,20 @@ void CoordinateManager::fill_arc_logical_to_translated() {
     }
 }
 
-std::shared_ptr<CoordinateManager> CoordinateManager::create_coordinate_manager(
-    tt::ARCH arch, const size_t tensix_harvesting_mask, const size_t dram_harvesting_mask) {
+void CoordinateManager::assert_create_coordinate_manager(
+    const tt::ARCH arch, const size_t tensix_harvesting_mask, const size_t dram_harvesting_mask) {
     log_assert(
         !(arch != tt::ARCH::BLACKHOLE && dram_harvesting_mask != 0), "DRAM harvesting is supported only for Blackhole");
 
+    if (arch == tt::ARCH::BLACKHOLE) {
+        log_assert(get_num_harvested(dram_harvesting_mask) <= 1, "Only one DRAM bank can be harvested on Blackhole");
+    }
+}
+
+std::shared_ptr<CoordinateManager> CoordinateManager::create_coordinate_manager(
+    tt::ARCH arch, const size_t tensix_harvesting_mask, const size_t dram_harvesting_mask) {
+    assert_create_coordinate_manager(arch, tensix_harvesting_mask, dram_harvesting_mask);
+
     switch (arch) {
         case tt::ARCH::GRAYSKULL:
             return create_coordinate_manager(
@@ -463,8 +472,7 @@ std::shared_ptr<CoordinateManager> CoordinateManager::create_coordinate_manager(
     const std::vector<tt_xy_pair>& arc_cores,
     const tt_xy_pair& pcie_grid_size,
     const std::vector<tt_xy_pair>& pcie_cores) {
-    log_assert(
-        !(arch != tt::ARCH::BLACKHOLE && dram_harvesting_mask != 0), "DRAM harvesting is supported only for Blackhole");
+    assert_create_coordinate_manager(arch, tensix_harvesting_mask, dram_harvesting_mask);
 
     switch (arch) {
         case tt::ARCH::GRAYSKULL:
@@ -527,3 +535,18 @@ size_t CoordinateManager::get_num_harvested(const size_t harvesting_mask) {
     }
     return num_harvested;
 }
+
+std::vector<size_t> CoordinateManager::get_harvested_indices(const size_t harvesting_mask) {
+    std::vector<size_t> indices;
+    size_t mask = harvesting_mask;
+    size_t index = 0;
+    while (mask > 0) {
+        if (mask & 1) {
+            indices.push_back(index);
+        }
+        mask >>= 1;
+        index++;
+    }
+
+    return indices;
+}
diff --git a/tests/api/test_core_coord_translation_bh.cpp b/tests/api/test_core_coord_translation_bh.cpp
@@ -252,6 +252,10 @@ TEST(CoordinateManager, CoordinateManagerBlackholeDRAMLogicalPhysicalMapping) {
     const std::vector<tt_xy_pair>& dram_cores = tt::umd::blackhole::DRAM_CORES;
 
     for (size_t harvesting_mask = 0; harvesting_mask < (1 << max_num_banks_harvested); harvesting_mask++) {
+        if (CoordinateManager::get_num_harvested(harvesting_mask) > 1) {
+            continue;
+        }
+
         std::shared_ptr<CoordinateManager> coordinate_manager =
             CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, harvesting_mask);
 
@@ -262,7 +266,7 @@ TEST(CoordinateManager, CoordinateManagerBlackholeDRAMLogicalPhysicalMapping) {
 
         for (size_t x = 0; x < num_dram_banks - num_banks_harvested; x++) {
             for (size_t y = 0; y < num_noc_ports_per_bank; y++) {
-                const CoreCoord logical_coords = CoreCoord(x, y, CoreType::TENSIX, CoordSystem::LOGICAL);
+                const CoreCoord logical_coords = CoreCoord(x, y, CoreType::DRAM, CoordSystem::LOGICAL);
                 const CoreCoord physical_coords = coordinate_manager->to(logical_coords, CoordSystem::PHYSICAL);
                 logical_to_physical[logical_coords] = physical_coords;
 
@@ -295,6 +299,10 @@ TEST(CoordinateManager, CoordinateManagerBlackholeDRAMLogicalVirtualMapping) {
     const size_t num_noc_ports_per_bank = tt::umd::blackhole::NUM_NOC_PORTS_PER_DRAM_BANK;
 
     for (size_t harvesting_mask = 0; harvesting_mask < (1 << max_num_banks_harvested); harvesting_mask++) {
+        if (CoordinateManager::get_num_harvested(harvesting_mask) > 1) {
+            continue;
+        }
+
         std::shared_ptr<CoordinateManager> coordinate_manager =
             CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, harvesting_mask);
 
@@ -305,7 +313,7 @@ TEST(CoordinateManager, CoordinateManagerBlackholeDRAMLogicalVirtualMapping) {
 
         for (size_t x = 0; x < num_dram_banks - num_harvested_banks; x++) {
             for (size_t y = 0; y < num_noc_ports_per_bank; y++) {
-                CoreCoord logical_coords = CoreCoord(x, y, CoreType::TENSIX, CoordSystem::LOGICAL);
+                CoreCoord logical_coords = CoreCoord(x, y, CoreType::DRAM, CoordSystem::LOGICAL);
                 CoreCoord virtual_coords = coordinate_manager->to(logical_coords, CoordSystem::VIRTUAL);
                 logical_to_virtual[logical_coords] = virtual_coords;
 
@@ -326,29 +334,66 @@ TEST(CoordinateManager, CoordinateManagerBlackholeDRAMLogicalVirtualMapping) {
     }
 }
 
-// Test equality of physical and translated coordinates for all logical coordinates for any harvesting mask.
-TEST(CoordinateManager, CoordinateManagerBlackholeDRAMPhysicalTranslatedEquality) {
+// Test DRAM translated mapping.
+TEST(CoordinateManager, CoordinateManagerBlackholeDRAMTranslatedMapping) {
     const size_t max_num_banks_harvested = tt::umd::blackhole::NUM_DRAM_BANKS;
     const size_t num_dram_banks = tt::umd::blackhole::NUM_DRAM_BANKS;
     const size_t num_noc_ports_per_bank = tt::umd::blackhole::NUM_NOC_PORTS_PER_DRAM_BANK;
 
     for (size_t harvesting_mask = 0; harvesting_mask < (1 << max_num_banks_harvested); harvesting_mask++) {
+        if (CoordinateManager::get_num_harvested(harvesting_mask) > 1) {
+            continue;
+        }
+
         std::shared_ptr<CoordinateManager> coordinate_manager =
             CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, harvesting_mask);
 
+        std::map<CoreCoord, CoreCoord> logical_to_translated;
+        std::set<CoreCoord> translated_coord_set;
+
         const size_t num_harvested_banks = CoordinateManager::get_num_harvested(harvesting_mask);
 
         for (size_t x = 0; x < num_dram_banks - num_harvested_banks; x++) {
             for (size_t y = 0; y < num_noc_ports_per_bank; y++) {
-                const CoreCoord logical_coords = CoreCoord(x, y, CoreType::TENSIX, CoordSystem::LOGICAL);
-
-                const CoreCoord physical_coords = coordinate_manager->to(logical_coords, CoordSystem::PHYSICAL);
+                const CoreCoord logical_coords = CoreCoord(x, y, CoreType::DRAM, CoordSystem::LOGICAL);
                 const CoreCoord translated_coords = coordinate_manager->to(logical_coords, CoordSystem::TRANSLATED);
 
-                EXPECT_EQ(physical_coords.x, translated_coords.x);
-                EXPECT_EQ(physical_coords.y, translated_coords.y);
+                EXPECT_GE(translated_coords.x, tt::umd::blackhole::dram_translated_coordinate_start_x);
+                EXPECT_GE(translated_coords.y, tt::umd::blackhole::dram_translated_coordinate_start_y);
+
+                logical_to_translated[logical_coords] = translated_coords;
+
+                // Expect that logical to translated translation is 1-1 mapping. No duplicates for translated
+                // coordinates.
+                EXPECT_EQ(translated_coord_set.count(translated_coords), 0);
+                translated_coord_set.insert(translated_coords);
             }
         }
+
+        for (auto it : logical_to_translated) {
+            const CoreCoord translated_coords = it.second;
+            const CoreCoord logical_coords = coordinate_manager->to(translated_coords, CoordSystem::LOGICAL);
+
+            // Expect that reverse mapping of translated coordinates gives the same logical coordinates
+            // using which we got the translated coordinates.
+            EXPECT_EQ(it.first, logical_coords);
+        }
+    }
+}
+
+// Test that we cannot create a coordinate manager with more than one DRAM bank harvested.
+TEST(CoordinateManager, CoordinateManagerBlackholeDRAMPMoreThanOneDRAMBankHarvested) {
+    const size_t max_num_banks_harvested = tt::umd::blackhole::NUM_DRAM_BANKS;
+    const size_t num_dram_banks = tt::umd::blackhole::NUM_DRAM_BANKS;
+    const size_t num_noc_ports_per_bank = tt::umd::blackhole::NUM_NOC_PORTS_PER_DRAM_BANK;
+
+    for (size_t harvesting_mask = 0; harvesting_mask < (1 << max_num_banks_harvested); harvesting_mask++) {
+        if (CoordinateManager::get_num_harvested(harvesting_mask) <= 1) {
+            continue;
+        }
+
+        EXPECT_THROW(
+            CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, harvesting_mask), std::runtime_error);
     }
 }