Enable CNN ops on blackhole

#0: skip PCC validation
tenstorrent · Nov 14, 2024 · 0ee3a89 · 0ee3a89
1 parent 52742eb
commit 0ee3a89
Show file tree

Hide file tree

Showing 8 changed files with 1,260 additions and 12 deletions.
diff --git a/models/demos/ttnn_resnet/tests/resnet50_test_infra.py b/models/demos/ttnn_resnet/tests/resnet50_test_infra.py
@@ -15,12 +15,17 @@
 from models.utility_functions import (
     is_wormhole_b0,
     is_grayskull,
+    is_blackhole,
     divup,
 )
 
 from tests.ttnn.utils_for_testing import assert_with_pcc
 from models.demos.ttnn_resnet.tt.custom_preprocessing import create_custom_mesh_preprocessor
-from models.demos.ttnn_resnet.tt.ttnn_functional_resnet50_new_conv_api import resnet50
+
+if is_blackhole():
+    from models.demos.ttnn_resnet.tt.ttnn_functional_resnet50 import resnet50
+else:
+    from models.demos.ttnn_resnet.tt.ttnn_functional_resnet50_new_conv_api import resnet50
 
 
 def load_resnet50_model(model_location_generator):
@@ -142,6 +147,7 @@ def load_resnet50_model(model_location_generator):
 }
 
 golden_pcc = {
+    ttnn.device.Arch.BLACKHOLE: golden_pcc,
     ttnn.device.Arch.WORMHOLE_B0: golden_pcc,
     ttnn.device.Arch.GRAYSKULL: golden_pcc,
 }
@@ -254,7 +260,7 @@ def setup_l1_sharded_input(self, device, torch_input_tensor=None):
         elif self.batch_size == 20:
             if is_grayskull():
                 core_grid = ttnn.CoreGrid(y=8, x=10)
-            elif is_wormhole_b0():
+            elif is_wormhole_b0() or is_blackhole():
                 core_grid = ttnn.CoreGrid(y=5, x=6)  # untested due to unsupported batch20 on WH
         num_devices = 1 if isinstance(device, ttnn.Device) else device.get_num_devices()
         # torch tensor