Wait for chip to be ready

tenstorrent · Feb 6, 2025 · 4b8a319 · 4b8a319
1 parent aa44485
commit 4b8a319
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 22 deletions.
diff --git a/device/api/umd/device/chip/chip.h b/device/api/umd/device/chip/chip.h
@@ -33,8 +33,6 @@ class Chip {
 
     const ChipInfo& get_chip_info();
 
-    virtual void wait_eth_cores_training(const uint32_t timeout_per_core_ms = 1000);
-
     // TODO: This should be private, once enough stuff is moved inside chip.
     // Probably also moved to LocalChip.
     tt_device_dram_address_params dram_address_params;
@@ -49,6 +47,11 @@ class Chip {
     ChipInfo chip_info_;
 
     tt_SocDescriptor soc_descriptor_;
+
+protected:
+    void wait_chip_to_be_ready();
+
+    virtual void wait_eth_cores_training(const uint32_t timeout_per_core_ms = 1000);
 };
 
 }  // namespace tt::umd
diff --git a/device/api/umd/device/chip/local_chip.h b/device/api/umd/device/chip/local_chip.h
@@ -20,11 +20,14 @@ class LocalChip : public Chip {
 
     bool is_mmio_capable() const override;
 
-    void wait_eth_cores_training(const uint32_t timeout_per_core_ms = 1000) override;
-
 private:
     std::unique_ptr<TTDevice> tt_device_;
 
+    void initialize_local_chip();
+
     void initialize_tlb_manager();
+
+protected:
+    void wait_eth_cores_training(const uint32_t timeout_per_core_ms = 1000) override;
 };
 }  // namespace tt::umd
diff --git a/device/chip/chip.cpp b/device/chip/chip.cpp
@@ -51,8 +51,8 @@ void Chip::set_barrier_address_params(const barrier_address_params& barrier_addr
 
 const ChipInfo& Chip::get_chip_info() { return chip_info_; }
 
-void Chip::wait_eth_cores_training(const uint32_t timeout_per_core) {
-    throw std::runtime_error("Waiting for training of ETH cores is supported only for Blackhole LocalChip.");
-}
+void Chip::wait_chip_to_be_ready() { wait_eth_cores_training(); }
+
+void Chip::wait_eth_cores_training(const uint32_t timeout_per_core) {}
 
 }  // namespace tt::umd
diff --git a/device/chip/local_chip.cpp b/device/chip/local_chip.cpp
@@ -15,19 +15,7 @@ namespace tt::umd {
 
 LocalChip::LocalChip(tt_SocDescriptor soc_descriptor, int pci_device_id) :
     Chip(soc_descriptor), tt_device_(TTDevice::create(pci_device_id)) {
-    initialize_tlb_manager();
-}
-
-void LocalChip::initialize_tlb_manager() {
-    auto tlb_manager = tt_device_->get_tlb_manager();
-    // Setup default dynamic tlbs.
-    tlb_manager->set_dynamic_tlb_config(
-        "LARGE_READ_TLB", tt_device_->get_architecture_implementation()->get_mem_large_read_tlb());
-    tlb_manager->set_dynamic_tlb_config(
-        "LARGE_WRITE_TLB", tt_device_->get_architecture_implementation()->get_mem_large_write_tlb());
-    tlb_manager->set_dynamic_tlb_config("REG_TLB", tt_device_->get_architecture_implementation()->get_reg_tlb());
-    tlb_manager->set_dynamic_tlb_config(
-        "SMALL_READ_WRITE_TLB", tt_device_->get_architecture_implementation()->get_small_read_write_tlb());
+    initialize_local_chip();
 }
 
 LocalChip::LocalChip(std::unique_ptr<TTDevice> tt_device) :
@@ -41,7 +29,24 @@ LocalChip::LocalChip(std::unique_ptr<TTDevice> tt_device) :
             tt_device->get_chip_info().noc_translation_enabled,
             tt_device->get_chip_info().harvesting_masks)),
     tt_device_(std::move(tt_device)) {
+    initialize_local_chip();
+}
+
+void LocalChip::initialize_local_chip() {
     initialize_tlb_manager();
+    wait_chip_to_be_ready();
+}
+
+void LocalChip::initialize_tlb_manager() {
+    auto tlb_manager = tt_device_->get_tlb_manager();
+    // Setup default dynamic tlbs.
+    tlb_manager->set_dynamic_tlb_config(
+        "LARGE_READ_TLB", tt_device_->get_architecture_implementation()->get_mem_large_read_tlb());
+    tlb_manager->set_dynamic_tlb_config(
+        "LARGE_WRITE_TLB", tt_device_->get_architecture_implementation()->get_mem_large_write_tlb());
+    tlb_manager->set_dynamic_tlb_config("REG_TLB", tt_device_->get_architecture_implementation()->get_reg_tlb());
+    tlb_manager->set_dynamic_tlb_config(
+        "SMALL_READ_WRITE_TLB", tt_device_->get_architecture_implementation()->get_small_read_write_tlb());
 }
 
 TTDevice* LocalChip::get_tt_device() { return tt_device_.get(); }

diff --git a/device/cluster.cpp b/device/cluster.cpp
@@ -3494,8 +3494,6 @@ std::unique_ptr<tt_ClusterDescriptor> Cluster::create_cluster_descriptor(
         const chip_id_t chip_id = it.first;
         const std::unique_ptr<Chip>& chip = it.second;
 
-        chip->wait_eth_cores_training();
-
         desc->all_chips.insert(chip_id);
         desc->chip_arch.insert({chip_id, chip->get_tt_device()->get_arch()});