Implement TTDevice get_board_type

tenstorrent · Mar 7, 2025 · 6270a14 · 6270a14
1 parent 7ae3cb1
commit 6270a14
Show file tree

Hide file tree

Showing 9 changed files with 1,005 additions and 935 deletions.
diff --git a/device/api/umd/device/tt_device/blackhole_tt_device.h b/device/api/umd/device/tt_device/blackhole_tt_device.h
@@ -25,6 +25,8 @@ class BlackholeTTDevice : public TTDevice {
 
     uint32_t get_clock() override;
 
+    BoardType get_board_type() override;
+
 private:
     static constexpr uint64_t ATU_OFFSET_IN_BH_BAR2 = 0x1200;
     std::set<size_t> iatu_regions_;

diff --git a/device/api/umd/device/tt_device/tt_device.h b/device/api/umd/device/tt_device/tt_device.h
@@ -141,6 +141,8 @@ class TTDevice {
 
     virtual uint32_t get_clock();
 
+    virtual BoardType get_board_type();
+
 protected:
     std::unique_ptr<PCIDevice> pci_device_;
     std::unique_ptr<architecture_implementation> architecture_impl_;

diff --git a/device/api/umd/device/tt_device/wormhole_tt_device.h b/device/api/umd/device/tt_device/wormhole_tt_device.h
@@ -16,5 +16,7 @@ class WormholeTTDevice : public TTDevice {
     ChipInfo get_chip_info() override;
 
     uint32_t get_clock() override;
+
+    BoardType get_board_type() override;
 };
 }  // namespace tt::umd
diff --git a/device/api/umd/device/types/cluster_descriptor_types.h b/device/api/umd/device/types/cluster_descriptor_types.h
@@ -122,6 +122,12 @@ inline BoardType get_board_type_from_board_id(const uint64_t board_id) {
         return BoardType::P100;
     } else if (upi == 0x40 || upi == 0x41 || upi == 0x42) {
         return BoardType::P150;
+    } else if (upi == 0x14) {
+        return BoardType::N300;
+    } else if (upi == 0x18) {
+        return BoardType::N150;
+    } else if (upi == 0xB) {
+        return BoardType::GALAXY;
     }
 
     throw std::runtime_error(fmt::format("No existing board type for board id {}", board_id));

diff --git a/device/tt_device/blackhole_tt_device.cpp b/device/tt_device/blackhole_tt_device.cpp
@@ -150,4 +150,10 @@ uint32_t BlackholeTTDevice::get_clock() {
     throw std::runtime_error("AICLK telemetry not available for Blackhole device.");
 }
 
+BoardType BlackholeTTDevice::get_board_type() {
+    return get_board_type_from_board_id(
+        ((uint64_t)telemetry->read_entry(blackhole::TAG_BOARD_ID_HIGH) << 32) |
+        (telemetry->read_entry(blackhole::TAG_BOARD_ID_LOW)));
+}
+
 }  // namespace tt::umd
diff --git a/device/tt_device/tt_device.cpp b/device/tt_device/tt_device.cpp
@@ -379,4 +379,10 @@ uint32_t TTDevice::get_clock() {
         "TTDevice is deleted.");
 }
 
+BoardType TTDevice::get_board_type() {
+    throw std::runtime_error(
+        "Base TTDevice class does not have get_board_type implemented. Move this to abstract function once Grayskull "
+        "TTDevice is deleted.");
+}
+
 }  // namespace tt::umd
diff --git a/device/tt_device/wormhole_tt_device.cpp b/device/tt_device/wormhole_tt_device.cpp
@@ -3,6 +3,9 @@
 // SPDX-License-Identifier: Apache-2.0
 #include "umd/device/tt_device/wormhole_tt_device.h"
 
+#include <iostream>
+#include <vector>
+
 #include "umd/device/wormhole_implementation.h"
 
 namespace tt::umd {
@@ -30,4 +33,22 @@ uint32_t WormholeTTDevice::get_clock() {
     return arc_msg_return_values[0];
 }
 
+BoardType WormholeTTDevice::get_board_type() {
+    ::std::vector<uint32_t> arc_msg_return_values = {0};
+    static const uint32_t timeout_ms = 1000;
+    uint32_t exit_code = get_arc_messenger()->send_message(
+        tt::umd::wormhole::ARC_MSG_COMMON_PREFIX | 0x2C, arc_msg_return_values, 0, 0, timeout_ms);
+
+    static constexpr uint64_t noc_telemetry_offset = 0x810000000;
+    uint64_t telemetry_struct_offset = arc_msg_return_values[0] + noc_telemetry_offset;
+
+    uint32_t board_id_lo;
+    uint32_t board_id_hi;
+    tt_xy_pair arc_core = tt::umd::wormhole::ARC_CORES[0];
+    read_from_device(&board_id_hi, arc_core, telemetry_struct_offset + 16, sizeof(uint32_t));
+    read_from_device(&board_id_lo, arc_core, telemetry_struct_offset + 20, sizeof(uint32_t));
+
+    return get_board_type_from_board_id(((uint64_t)board_id_hi << 32) | board_id_lo);
+}
+
 }  // namespace tt::umd
diff --git a/tests/api/test_tt_device.cpp b/tests/api/test_tt_device.cpp
@@ -23,7 +23,7 @@ tt_xy_pair get_any_tensix_core(tt::ARCH arch) {
     }
 }
 
-TEST(TTDeviceTest, BasicTTDeviceIO) {
+TEST(ApiTTDeviceTest, BasicTTDeviceIO) {
     std::vector<int> pci_device_ids = PCIDevice::enumerate_devices();
 
     uint64_t address = l1_mem::address_map::NCRISC_FIRMWARE_BASE;
@@ -44,3 +44,16 @@ TEST(TTDeviceTest, BasicTTDeviceIO) {
         data_read = std::vector<uint32_t>(data_write.size(), 0);
     }
 }
+
+TEST(ApiTTDeviceTest, TTDeviceGetBoardType) {
+    std::vector<int> pci_device_ids = PCIDevice::enumerate_devices();
+    for (int pci_device_id : pci_device_ids) {
+        std::unique_ptr<TTDevice> tt_device = TTDevice::create(pci_device_id);
+
+        BoardType board_type = tt_device->get_board_type();
+
+        EXPECT_TRUE(
+            board_type == BoardType::N150 || board_type == BoardType::N300 || board_type == BoardType::P100 ||
+            board_type == BoardType::P150 || board_type == BoardType::P300 || board_type == BoardType::GALAXY);
+    }
+}