addressed some comments

tenstorrent · Dec 5, 2024 · 0f466e7 · 0f466e7
1 parent f27b7bf
commit 0f466e7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/device/api/umd/device/cluster.h b/device/api/umd/device/cluster.h
@@ -819,7 +819,7 @@ class Cluster : public tt_device {
     virtual std::uint32_t get_numa_node_for_pcie_device(std::uint32_t device_id);
     virtual tt_version get_ethernet_fw_version() const;
 
-    TTDevice* get_tt_device(int device_id) const;
+    TTDevice* get_tt_device(chip_id_t device_id) const;
 
     // Destructor
     virtual ~Cluster();
@@ -973,7 +973,9 @@ class Cluster : public tt_device {
     std::set<chip_id_t> target_devices_in_cluster = {};
     std::set<chip_id_t> target_remote_chips = {};
     tt::ARCH arch_name;
-    std::unordered_map<chip_id_t, std::unique_ptr<TTDevice>> m_tt_device_map;  // Map of enabled tt devices
+
+    // Map of enabled tt devices
+    std::unordered_map<chip_id_t, std::unique_ptr<TTDevice>> m_tt_device_map;
     std::shared_ptr<tt_ClusterDescriptor> cluster_desc;
 
     // remote eth transfer setup

diff --git a/device/cluster.cpp b/device/cluster.cpp
@@ -1730,7 +1730,7 @@ void* Cluster::host_dma_address(std::uint64_t offset, chip_id_t src_device_id, u
 }
 
 // Wrapper for throwing more helpful exception when not-enabled pci intf is accessed.
-inline TTDevice* Cluster::get_tt_device(int device_id) const {
+inline TTDevice* Cluster::get_tt_device(chip_id_t device_id) const {
     if (!m_tt_device_map.count(device_id)) {
         throw std::runtime_error(fmt::format("device_id: {} attempted to be accessed, but is not enabled.", device_id));
     }