From 3be2fd5efe4b8d9a82c994a3058e5dcce64dbcbe Mon Sep 17 00:00:00 2001
From: Yung-Ti Chou <153475060+swimdi@users.noreply.github.com>
Date: Mon, 30 Dec 2024 16:49:02 +0800
Subject: [PATCH] =?UTF-8?q?Sweep=20the=20=E2=9D=8C=20models=20again(no=20t?=
 =?UTF-8?q?rain)=20(#690)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* sweep failed model

* Regist converted_end_to_end mark

* Skip calculate retinanet_resnet50_fpn_v2 because before/after shape not the same
---
 tests/conftest.py                             |  2 +-
 .../test_torchvision_object_detection.py      |  6 ++---
 tests/pytest.ini                              |  1 +
 torch_ttnn/passes/lowering/to_tt_guard.py     |  6 +++--
 .../passes/lowering/to_tt_guard_autogen.py    | 22 +++++++++++++++++++
 torch_ttnn/passes/lowering/to_tt_pass.py      |  4 +++-
 6 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 0416ba56b..e414ec8a8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -188,7 +188,7 @@ def compile_and_run(device, reset_torch_dynamo, request):
             if len(option._out_fx_graphs) > 0:
                 option._out_fx_graphs[0].print_tabular()
 
-            if model_name not in ["speecht5-tts", "ssd300_vgg16"]:
+            if model_name not in ["speecht5-tts", "ssd300_vgg16", "retinanet_resnet50_fpn_v2"]:
                 accuracy = calculate_accuracy(outputs, outputs_after)
                 if accuracy:
                     comp_runtime_metrics["accuracy"] = accuracy
diff --git a/tests/models/torchvision/test_torchvision_object_detection.py b/tests/models/torchvision/test_torchvision_object_detection.py
index 8e42af73c..aa4d998cb 100644
--- a/tests/models/torchvision/test_torchvision_object_detection.py
+++ b/tests/models/torchvision/test_torchvision_object_detection.py
@@ -42,10 +42,8 @@ def _load_inputs(self):
     [
         ("ssd300_vgg16", "SSD300_VGG16_Weights"),
         ("ssdlite320_mobilenet_v3_large", "SSDLite320_MobileNet_V3_Large_Weights"),
-        pytest.param(("retinanet_resnet50_fpn", "RetinaNet_ResNet50_FPN_Weights"), marks=pytest.mark.compilation_xfail),
-        pytest.param(
-            ("retinanet_resnet50_fpn_v2", "RetinaNet_ResNet50_FPN_V2_Weights"), marks=pytest.mark.compilation_xfail
-        ),
+        ("retinanet_resnet50_fpn", "RetinaNet_ResNet50_FPN_Weights"),
+        ("retinanet_resnet50_fpn_v2", "RetinaNet_ResNet50_FPN_V2_Weights"),
     ],
 )
 def test_torchvision_object_detection(record_property, model_info, mode):
diff --git a/tests/pytest.ini b/tests/pytest.ini
index 1b79baf5c..8696260b8 100644
--- a/tests/pytest.ini
+++ b/tests/pytest.ini
@@ -2,3 +2,4 @@
 markers =
     compilation_xfail: marks tests with compiled run as xfail but does not change torch run
     skip_platform: marks tests that are not compatible with specified platform
+    converted_end_to_end: marks tests that all aten ops are converted to ttnn ops
diff --git a/torch_ttnn/passes/lowering/to_tt_guard.py b/torch_ttnn/passes/lowering/to_tt_guard.py
index 42868e7ac..bfe128809 100644
--- a/torch_ttnn/passes/lowering/to_tt_guard.py
+++ b/torch_ttnn/passes/lowering/to_tt_guard.py
@@ -79,8 +79,9 @@
 # EXTRA BLOCKLIST OF retinanet_resnet50_fpn
 ############################################################
 # Statically allocated circular buffers on core range [(x=0,y=0) - (x=0,y=0)] grow to 3580704 B which is beyond max L1 size of 1499136 B
-# TODO: not pass yet
-
+aten_aten_stack_default = [
+    ["List[Tensor] tensors = [<[13600]>, <[13600]>, <[13600]>, <[13600]>]", "int dim = 1"],
+]
 
 ############################################################
 # EXTRA BLOCKLIST OF retinanet_resnet50_fpn_v2
@@ -92,6 +93,7 @@
 
 GUARD[torch.ops.aten.gt.Scalar] = partial(guard_aten, aten_gt_Scalar_blocklist)
 GUARD[torch.ops.aten.cumsum.default] = partial(guard_aten, aten_cumsum_default_blocklist)
+GUARD[torch.ops.aten.stack.default] = partial(guard_aten, aten_aten_stack_default)
 
 
 def can_lowering_to_ttnn(node):
diff --git a/torch_ttnn/passes/lowering/to_tt_guard_autogen.py b/torch_ttnn/passes/lowering/to_tt_guard_autogen.py
index 148defbe5..75668a40b 100644
--- a/torch_ttnn/passes/lowering/to_tt_guard_autogen.py
+++ b/torch_ttnn/passes/lowering/to_tt_guard_autogen.py
@@ -286,6 +286,28 @@
         "List[int] output_padding = [0, 0]",
         "int groups = 960",
     ],
+    [
+        "Tensor<[1, 768, 3000]> input = ?",
+        "Tensor<[768, 768, 3]> weight = ?",
+        "Optional[Tensor]<[768]> bias = ?",
+        "List[int] stride = [2]",
+        "List[int] padding = [1]",
+        "List[int] dilation = [1]",
+        "bool transposed = False",
+        "List[int] output_padding = [0]",
+        "int groups = 1",
+    ],
+    [
+        "Tensor<[1, 320, 64, 64]> input = ?",
+        "Tensor<[320, 320, 3, 3]> weight = ?",
+        "Optional[Tensor]<[320]> bias = ?",
+        "List[int] stride = [1, 1]",
+        "List[int] padding = [1, 1]",
+        "List[int] dilation = [1, 1]",
+        "bool transposed = False",
+        "List[int] output_padding = [0, 0]",
+        "int groups = 1",
+    ],
     # TODO(tt-metal#16173): weight_matrix_width_ntiles % weight_block_w_ntiles == 0
     [
         "Tensor<[1, 1232, 14, 14]> input = ?",
diff --git a/torch_ttnn/passes/lowering/to_tt_pass.py b/torch_ttnn/passes/lowering/to_tt_pass.py
index b3e9b6017..a3c91260d 100644
--- a/torch_ttnn/passes/lowering/to_tt_pass.py
+++ b/torch_ttnn/passes/lowering/to_tt_pass.py
@@ -706,7 +706,9 @@ def reshape_1d(code, args=args, kwargs=kwargs):
 
             if node.target == torch.ops.aten.repeat.default:
                 tensor, sizes = args
-                shape = tensor.meta["val"].size()
+                shape = get_shape(gm, tensor)
+                if shape is None:
+                    return None
 
                 if np.prod(sizes) == 1:
                     return tensor