fix gs tests

tenstorrent · Dec 17, 2024 · eb8f6a8 · eb8f6a8
1 parent 9a47a2f
commit eb8f6a8
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/grayskull/test_cluster_gs.cpp b/tests/grayskull/test_cluster_gs.cpp
@@ -9,9 +9,9 @@
 #include "tests/test_utils/device_test_utils.hpp"
 #include "tests/test_utils/generate_cluster_desc.hpp"
 #include "umd/device/cluster.h"
+#include "umd/device/grayskull_implementation.h"
 #include "umd/device/tt_cluster_descriptor.h"
 #include "umd/device/tt_soc_descriptor.h"
-#include "umd/device/wormhole_implementation.h"
 
 using namespace tt::umd;
 
@@ -95,7 +95,7 @@ TEST(SiliconDriverGS, CustomSocDesc) {
 
 TEST(SiliconDriverGS, HarvestingRuntime) {
     auto get_static_tlb_index = [](tt_xy_pair target) {
-        int flat_index = target.y * tt::umd::wormhole::GRID_SIZE_X + target.x;
+        int flat_index = target.y * tt::umd::grayskull::GRID_SIZE_X + target.x;
         if (flat_index == 0) {
             return -1;
         }
@@ -188,7 +188,7 @@ TEST(SiliconDriverGS, HarvestingRuntime) {
 
 TEST(SiliconDriverGS, StaticTLB_RW) {
     auto get_static_tlb_index = [](tt_xy_pair target) {
-        int flat_index = target.y * tt::umd::wormhole::GRID_SIZE_X + target.x;
+        int flat_index = target.y * tt::umd::grayskull::GRID_SIZE_X + target.x;
         if (flat_index == 0) {
             return -1;
         }
@@ -403,7 +403,7 @@ TEST(SiliconDriverGS, MultiThreadedMemBar) {  // this tests takes ~5 mins to run
     // Memory barrier flags get sent to address 0 for all channels in this test
 
     auto get_static_tlb_index = [](tt_xy_pair target) {
-        int flat_index = target.y * tt::umd::wormhole::GRID_SIZE_X + target.x;
+        int flat_index = target.y * tt::umd::grayskull::GRID_SIZE_X + target.x;
         if (flat_index == 0) {
             return -1;
         }