From 6487de1a02eb4ed8ddfc342f288c95abf50b1c18 Mon Sep 17 00:00:00 2001
From: pjanevski <pjanevski@tenstorrent.com>
Date: Thu, 26 Dec 2024 15:40:18 +0000
Subject: [PATCH] Address comments

---
 device/api/umd/device/blackhole_implementation.h | 11 ++++++++---
 device/blackhole/blackhole_implementation.cpp    | 13 +++++++++----
 tests/api/test_core_coord_translation_bh.cpp     |  4 ++--
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/device/api/umd/device/blackhole_implementation.h b/device/api/umd/device/blackhole_implementation.h
index 370060de..f5047c9e 100644
--- a/device/api/umd/device/blackhole_implementation.h
+++ b/device/api/umd/device/blackhole_implementation.h
@@ -100,9 +100,9 @@ static const std::vector<tt_xy_pair> ARC_CORES = {{8, 0}};
 static const std::vector<tt_xy_pair> ARC_LOCATIONS = ARC_CORES;
 
 static const tt_xy_pair PCIE_GRID_SIZE = {1, 1};
-static const std::vector<tt_xy_pair> PCIE_CORES_LOCAL = {{{2, 0}}};
-static const std::vector<tt_xy_pair> PCI_LOCATIONS = PCIE_CORES_LOCAL;
-static const std::vector<tt_xy_pair> PCIE_CORES_REMOTE = {{{11, 0}}};
+static const std::vector<tt_xy_pair> PCIE_CORES_TYPE2 = {{{2, 0}}};
+static const std::vector<tt_xy_pair> PCI_LOCATIONS = PCIE_CORES_TYPE2;
+static const std::vector<tt_xy_pair> PCIE_CORES_TYPE1 = {{{11, 0}}};
 
 static const tt_xy_pair ETH_GRID_SIZE = {14, 1};
 static const std::vector<tt_xy_pair> ETH_CORES = {
@@ -197,6 +197,11 @@ 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;
 
+/*
+ * Ge the PCIE core that can be used for communication with host
+ * based on the board type and whether the chip is remote or not.
+ * Information on remote chip is used only if the board type is P300.
+ */
 std::vector<tt_xy_pair> get_pcie_cores(const BoardType board_type, const bool is_chip_remote);
 
 }  // namespace blackhole
diff --git a/device/blackhole/blackhole_implementation.cpp b/device/blackhole/blackhole_implementation.cpp
index be33cf0f..1c941249 100644
--- a/device/blackhole/blackhole_implementation.cpp
+++ b/device/blackhole/blackhole_implementation.cpp
@@ -8,6 +8,7 @@
 #include "blackhole/eth_l1_address_map.h"
 #include "blackhole/host_mem_address_map.h"
 #include "blackhole/l1_address_map.h"
+#include "logger.hpp"
 #include "umd/device/cluster.h"
 
 constexpr std::uint32_t NOC_ADDR_LOCAL_BITS = 36;   // source: noc_parameters.h, common for WH && BH
@@ -106,15 +107,19 @@ tt_driver_noc_params blackhole_implementation::get_noc_params() const {
 
 namespace blackhole {
 std::vector<tt_xy_pair> get_pcie_cores(const BoardType board_type, const bool is_chip_remote) {
+    if (is_chip_remote) {
+        log_assert(board_type == BoardType::P300, "Remote chip is supported only for Blackhole P300 board.");
+    }
+
     if (board_type == BoardType::UNKNOWN || board_type == BoardType::P100) {
-        return PCIE_CORES_REMOTE;
+        return PCIE_CORES_TYPE1;
     } else if (board_type == BoardType::P150A) {
-        return PCIE_CORES_LOCAL;
+        return PCIE_CORES_TYPE2;
     } else if (board_type == BoardType::P300) {
-        return is_chip_remote ? PCIE_CORES_REMOTE : PCIE_CORES_LOCAL;
+        return is_chip_remote ? PCIE_CORES_TYPE1 : PCIE_CORES_TYPE2;
     }
 
-    return PCIE_CORES_LOCAL;
+    return PCIE_CORES_TYPE2;
 }
 }  // namespace blackhole
 
diff --git a/tests/api/test_core_coord_translation_bh.cpp b/tests/api/test_core_coord_translation_bh.cpp
index 7c002edb..994a4208 100644
--- a/tests/api/test_core_coord_translation_bh.cpp
+++ b/tests/api/test_core_coord_translation_bh.cpp
@@ -500,7 +500,7 @@ TEST(CoordinateManager, CoordinateManagerBlackholePCIETranslationLocal) {
     std::shared_ptr<CoordinateManager> coordinate_manager =
         CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, 0, 0, BoardType::P300, false);
     const tt_xy_pair pcie_grid_size = tt::umd::blackhole::PCIE_GRID_SIZE;
-    const std::vector<tt_xy_pair> pcie_cores = tt::umd::blackhole::PCIE_CORES_LOCAL;
+    const std::vector<tt_xy_pair> pcie_cores = tt::umd::blackhole::PCIE_CORES_TYPE2;
 
     for (size_t x = 0; x < pcie_grid_size.x; x++) {
         for (size_t y = 0; y < pcie_grid_size.y; y++) {
@@ -523,7 +523,7 @@ TEST(CoordinateManager, CoordinateManagerBlackholePCIETranslationRemote) {
     std::shared_ptr<CoordinateManager> coordinate_manager =
         CoordinateManager::create_coordinate_manager(tt::ARCH::BLACKHOLE, 0, 0, 0, BoardType::P300, true);
     const tt_xy_pair pcie_grid_size = tt::umd::blackhole::PCIE_GRID_SIZE;
-    const std::vector<tt_xy_pair> pcie_cores = tt::umd::blackhole::PCIE_CORES_REMOTE;
+    const std::vector<tt_xy_pair> pcie_cores = tt::umd::blackhole::PCIE_CORES_TYPE1;
 
     for (size_t x = 0; x < pcie_grid_size.x; x++) {
         for (size_t y = 0; y < pcie_grid_size.y; y++) {