Remove some ARCH_NAME ENV usage at runtime (#16825)

tenstorrent · Jan 17, 2025 · 582d79a · 582d79a
1 parent 422efe4
commit 582d79a
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 9 deletions.
diff --git a/tests/tt_metal/tt_metal/common/device_fixture.hpp b/tests/tt_metal/tt_metal/common/device_fixture.hpp
@@ -95,7 +95,7 @@ class BlackholeSingleCardFixture : public DeviceSingleCardFixture {
 protected:
     void SetUp() override {
         this->validate_dispatch_mode();
-        this->arch_ = tt::get_arch_from_string(tt::test_utils::get_env_arch_name());
+        this->arch_ = tt::get_arch_from_string(tt::test_utils::get_umd_arch_name());
         if (this->arch_ != tt::ARCH::BLACKHOLE) {
             GTEST_SKIP();
         }

diff --git a/...metal/tt_metal/dispatch/dispatch_buffer/test_EnqueueWriteBuffer_and_EnqueueReadBuffer.cpp b/...metal/tt_metal/dispatch/dispatch_buffer/test_EnqueueWriteBuffer_and_EnqueueReadBuffer.cpp
@@ -1421,8 +1421,7 @@ TEST_F(CommandQueueSingleCardBufferFixture, ShardedBufferLargeDRAMReadWrites) {
 }
 
 TEST_F(CommandQueueSingleCardBufferFixture, StressWrapTest) {
-    const char* arch = getenv("ARCH_NAME");
-    if (strcasecmp(arch, "wormhole_b0") == 0) {
+    if (this->arch_ == tt::ARCH::WORMHOLE_B0) {
         tt::log_info("cannot run this test on WH B0");
         GTEST_SKIP();
         return;  // skip for WH B0

diff --git a/tests/tt_metal/tt_metal/integration/matmul/test_matmul_multi_core_X_dram.cpp b/tests/tt_metal/tt_metal/integration/matmul/test_matmul_multi_core_X_dram.cpp
@@ -541,11 +541,10 @@ bool matmul_multi_core_multi_dram(DispatchFixture* fixture, tt_metal::IDevice* d
 }  // namespace unit_tests_common::matmul::test_matmul_multi_core_X_dram
 
 TEST_F(DispatchFixture, TensixMatmulMultiCoreSingleDRAM) {
-    const char* arch = getenv("ARCH_NAME");
     if (!getenv("TT_METAL_SLOW_DISPATCH_MODE")) {
         log_info(LogTest, "This test is only supported in slow dispatch mode");
         GTEST_SKIP();
-    } else if (strcasecmp(arch, "wormhole_b0") == 0) {
+    } else if (this->arch_ == tt::ARCH::WORMHOLE_B0) {
         tt::log_info("This test is disabled in WH B0");
         GTEST_SKIP();
     }

diff --git a/tests/tt_metal/tt_metal/perf_microbenchmark/dispatch/test_pgm_dispatch.cpp b/tests/tt_metal/tt_metal/perf_microbenchmark/dispatch/test_pgm_dispatch.cpp
@@ -4,6 +4,7 @@
 
 #include "umd/device/types/cluster_descriptor_types.h"
 #include <tt-metalium/host_api.hpp>
+#include <tt-metalium/hal_exp.hpp>
 #include <tt-metalium/tt_metal.hpp>
 #include <tt-metalium/command_queue.hpp>
 #include <tt-metalium/device.hpp>
@@ -55,7 +56,7 @@ std::tuple<uint32_t, uint32_t> get_core_count() {
     uint32_t core_x = 0;
     uint32_t core_y = 0;
 
-    std::string arch_name{getenv("ARCH_NAME")};
+    std::string arch_name = tt::tt_metal::experimental::hal::get_arch_name();
     if (arch_name == "grayskull") {
         core_x = 11;
         core_y = 8;
@@ -603,7 +604,8 @@ int main(int argc, char** argv) {
             .use_all_cores = true})
         ->Apply(Max8192Args)
         ->UseManualTime();
-    if (getenv("ARCH_NAME") == std::string("wormhole_b0")) {
+    std::string arch_name = tt::tt_metal::experimental::hal::get_arch_name();
+    if (arch_name == std::string("wormhole_b0")) {
         benchmark::RegisterBenchmark(
             "BM_pgm_dispatch/eth_dispatch",
             BM_pgm_dispatch,

diff --git a/tt-train/sources/examples/simple_cnn/main.cpp b/tt-train/sources/examples/simple_cnn/main.cpp
@@ -10,9 +10,7 @@ int main() {
     const size_t tensor_height = 32;
 
     std::srand(0);
-    auto arch_ = tt::get_arch_from_string(tt::test_utils::get_env_arch_name());
     auto num_devices_ = tt::tt_metal::GetNumAvailableDevices();
-    std::cout << "Arch:" << tt::test_utils::get_env_arch_name() << std::endl;
     std::cout << "num_devices:" << num_devices_ << std::endl;
     auto device = tt::tt_metal::CreateDevice(0);
     std::cout << "Device created" << std::endl;