Adds license comments

Author: ign-amanks

models/experimental/SSR/demo/ssr_demo.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import os
 import torch
 import ttnn

models/experimental/SSR/tests/common/test_mlp.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import torch
 import pytest
 
@@ -83,4 +86,4 @@ def test_mlp(device, in_features, hidden_features, out_features, input_shape):
     else:
         logger.warning("SSR MLP Failed!")
 
-    assert does_pass
+    assert does_pass, f"PCC check failed: {pcc_message}"

models/experimental/SSR/tests/test_ssr.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import torch
 import pytest
 import ttnn
@@ -12,10 +15,8 @@
 from models.experimental.SSR.tests.tile_refinement.test_HAB import create_relative_position_index
 
 from ttnn.model_preprocessing import preprocess_model_parameters
-from models.utility_functions import (
-    tt2torch_tensor,
-    comp_pcc,
-)
+from models.utility_functions import tt2torch_tensor
+from tests.ttnn.utils_for_testing import check_with_pcc
 
 
 def create_ssr_preprocessor(device, args, num_cls):
@@ -166,11 +167,8 @@ def test_ssr_model(input_shape, num_cls, with_conv):
         tt_torch_sr = tt_torch_sr.permute(0, 3, 1, 2)
 
         # Compare outputs
-        sr_pass, sr_pcc_message = comp_pcc(ref_sr, tt_torch_sr, 0.95)
-        fea3_pass, fea3_pcc_message = comp_pcc(ref_patch_fea3, tt_torch_patch_fea3, 0.95)
-
-        logger.info(f"SR Output PCC: {sr_pcc_message}")
-        logger.info(f"Patch Fea3 PCC: {fea3_pcc_message}")
+        sr_pass, sr_pcc_message = check_with_pcc(ref_sr, tt_torch_sr, 0.95)
+        fea3_pass, fea3_pcc_message = check_with_pcc(ref_patch_fea3, tt_torch_patch_fea3, 0.95)
 
         all_pass = sr_pass and fea3_pass
 
@@ -179,8 +177,8 @@ def test_ssr_model(input_shape, num_cls, with_conv):
         else:
             logger.warning("TTSSR Test Failed!")
 
-        assert sr_pass, f"SR output comparison failed: {sr_pcc_message}"
-        assert fea3_pass, f"Patch fea3 comparison failed: {fea3_pcc_message}"
+        assert sr_pass, f"SR output failed PCC check: {sr_pcc_message}"
+        assert fea3_pass, f"Patch fea3 failed PCC check: {fea3_pcc_message}"
 
     finally:
         ttnn.close_device(device)

models/experimental/SSR/tests/tile_refinement/test_CAB.py

@@ -118,20 +118,9 @@ def test_cab_block(device, batch_size, num_feat, height, width, compress_ratio,
     # Compare outputs
     does_pass, pcc_message = check_with_pcc(ref_output, tt_torch_output, 0.97)
 
-    logger.info(f"Batch: {batch_size}, Features: {num_feat}, Size: {height}x{width}")
-    logger.info(f"Compress ratio: {compress_ratio}, Squeeze factor: {squeeze_factor}")
-    logger.info(f"Reference output shape: {ref_output.shape}")
-    logger.info(f"TTNN output shape: {tt_torch_output.shape}")
-    logger.info(pcc_message)
-
     if does_pass:
         logger.info("CAB Block Passed!")
     else:
         logger.warning("CAB Block Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
-
-    # Verify output shapes match
-    assert (
-        ref_output.shape == tt_torch_output.shape
-    ), f"Shape mismatch: ref {ref_output.shape} vs ttnn {tt_torch_output.shape}"

models/experimental/SSR/tests/tile_refinement/test_HAB.py

@@ -313,19 +313,9 @@ def test_hab_block(device, batch_size, height, width, dim, num_heads, window_siz
     # Compare outputs
     does_pass, pcc_message = check_with_pcc(ref_output, tt_torch_output, 0.95)
 
-    logger.info(f"Batch: {batch_size}, Size: {height}x{width}, Dim: {dim}")
-    logger.info(f"Heads: {num_heads}, Window: {window_size}, Shift: {shift_size}")
-    logger.info(f"Reference output shape: {ref_output.shape}")
-    logger.info(f"TTNN output shape: {tt_torch_output.shape}")
-    logger.info(f"Actual Run: {actual_run_time} s")
-    logger.info(pcc_message)
-
     if does_pass:
         logger.info("HAB Block Passed!")
     else:
         logger.warning("HAB Block Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
-    assert (
-        ref_output.shape == tt_torch_output.shape
-    ), f"Shape mismatch: ref {ref_output.shape} vs ttnn {tt_torch_output.shape}"

models/experimental/SSR/tests/tile_refinement/test_OCAB.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import pytest
 import torch
 import torch.nn as nn
@@ -97,7 +100,8 @@ def custom_preprocessor(model, name):
     "dim, input_resolution, window_size, overlap_ratio, num_heads, input_shape",
     ((180, (64, 64), 16, 0.5, 6, (1, 4096, 180)),),
 )
-def test_ocab(dim, input_resolution, window_size, overlap_ratio, num_heads, input_shape):
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
+def test_ocab(device, dim, input_resolution, window_size, overlap_ratio, num_heads, input_shape):
     x = torch.randn(input_shape)
 
     # Create reference OCAB layer
@@ -121,41 +125,36 @@ def test_ocab(dim, input_resolution, window_size, overlap_ratio, num_heads, inpu
 
     ref_output = ref_layer(x, x_size, rpi)
 
-    device = ttnn.open_device(device_id=0, l1_small_size=32768)
     ttnn.synchronize_device(device)
 
-    try:
-        parameters = preprocess_model_parameters(
-            initialize_model=lambda: ref_layer, custom_preprocessor=create_ocab_preprocessor(device), device=device
-        )
-
-        tt_layer = TTOCAB(
-            device=device,
-            dim=dim,
-            input_resolution=input_resolution,
-            window_size=window_size,
-            overlap_ratio=overlap_ratio,
-            num_heads=num_heads,
-            parameters=parameters,
-        )
-
-        tt_input = ttnn.from_torch(
-            x, device=device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16, memory_config=ttnn.L1_MEMORY_CONFIG
-        )
-        tt_rpi = ttnn.from_torch(rpi, device=device, layout=ttnn.TILE_LAYOUT)
-        tt_output = tt_layer.forward(tt_input, x_size, tt_rpi)
-        tt_torch_output = tt2torch_tensor(tt_output)
-
-        does_pass, pcc_message = check_with_pcc(ref_output, tt_torch_output, 0.99)
-
-        logger.info(pcc_message)
-
-        if does_pass:
-            logger.info("OCAB Layer Passed!")
-        else:
-            logger.warning("OCAB Layer Failed!")
-
-    finally:
-        ttnn.close_device(device)
-
-    assert does_pass
+    parameters = preprocess_model_parameters(
+        initialize_model=lambda: ref_layer, custom_preprocessor=create_ocab_preprocessor(device), device=device
+    )
+
+    tt_layer = TTOCAB(
+        device=device,
+        dim=dim,
+        input_resolution=input_resolution,
+        window_size=window_size,
+        overlap_ratio=overlap_ratio,
+        num_heads=num_heads,
+        parameters=parameters,
+    )
+
+    tt_input = ttnn.from_torch(
+        x, device=device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16, memory_config=ttnn.L1_MEMORY_CONFIG
+    )
+    tt_rpi = ttnn.from_torch(rpi, device=device, layout=ttnn.TILE_LAYOUT)
+    tt_output = tt_layer.forward(tt_input, x_size, tt_rpi)
+    tt_torch_output = tt2torch_tensor(tt_output)
+
+    does_pass, pcc_message = check_with_pcc(ref_output, tt_torch_output, 0.99)
+
+    logger.info(pcc_message)
+
+    if does_pass:
+        logger.info("OCAB Layer Passed!")
+    else:
+        logger.warning("OCAB Layer Failed!")
+
+    assert does_pass, f"PCC check failed: {pcc_message}"

models/experimental/SSR/tests/tile_refinement/test_atten_blocks.py

@@ -136,19 +136,9 @@ def test_atten_blocks(device, batch_size, height, width, dim, num_heads, window_
     # Compare outputs
     does_pass, pcc_message = check_with_pcc(ref_output, tt_torch_output, 0.90)
 
-    logger.info(f"Batch: {batch_size}, Size: {height}x{width}, Dim: {dim}")
-    logger.info(f"Heads: {num_heads}, Window: {window_size}, Depth: {depth}")
-    logger.info(f"Overlap ratio: {overlap_ratio}, MLP ratio: {mlp_ratio}")
-    logger.info(f"Reference output shape: {ref_output.shape}")
-    logger.info(f"TTNN output shape: {tt_torch_output.shape}")
-    logger.info(pcc_message)
-
     if does_pass:
         logger.info("AttenBlocks Passed!")
     else:
         logger.warning("AttenBlocks Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
-    assert (
-        ref_output.shape == tt_torch_output.shape
-    ), f"Shape mismatch: ref {ref_output.shape} vs ttnn {tt_torch_output.shape}"

models/experimental/SSR/tests/tile_refinement/test_patch_embed_tile_refinement.py

@@ -161,5 +161,3 @@ def test_patch_embed_simple(device, batch_size, img_size, patch_size, in_chans,
         logger.warning("TR PatchEmbed Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
-
-    ttnn.close_device(device)

models/experimental/SSR/tests/tile_refinement/test_patch_unembed.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import pytest
 import torch
 import ttnn
@@ -57,5 +60,3 @@ def test_tt_patch_unembed(device, batch_size, img_size, patch_size, in_chans, em
         logger.warning("TR PatchEmbed Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
-
-    ttnn.close_device(device)

models/experimental/SSR/tests/tile_refinement/test_tile_refinement.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
 import torch
 import pytest
 import ttnn
@@ -118,8 +121,9 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
         (64, 2, 180, (6, 6, 6, 6, 6, 6), (6, 6, 6, 6, 6, 6), 16, 2, 4, (3, 3, 64, 64)),
     ],
 )
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
 def test_tile_refinement(
-    img_size, patch_size, embed_dim, depths, num_heads, window_size, mlp_ratio, upscale, input_shape
+    device, img_size, patch_size, embed_dim, depths, num_heads, window_size, mlp_ratio, upscale, input_shape
 ):
     """Test TTTileRefinement model against PyTorch reference"""
 
@@ -166,8 +170,6 @@ def test_tile_refinement(
     # Create params dictionary
     params = {"rpi_sa": rpi_sa, "attn_mask": attn_mask, "rpi_oca": rpi_oca}
 
-    device = ttnn.open_device(device_id=0, l1_small_size=32768)
-
     tt_rpi_sa = ttnn.from_torch(rpi_sa, device=device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=ttnn.uint32)
 
     tt_rpi_oca = ttnn.from_torch(rpi_oca, device=device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.uint32)
@@ -179,11 +181,7 @@ def test_tile_refinement(
     # Get reference output (both image and features)
     with torch.no_grad():
         ref_output, ref_features = ref_model(x)
-        # ref_output = ref_model(x)
 
-    # Open TTNN device
-
-    try:
         # Preprocess model parameters
         parameters = preprocess_model_parameters(
             initialize_model=lambda: ref_model,
@@ -247,6 +245,3 @@ def test_tile_refinement(
 
         assert output_pass, f"Output comparison failed: {output_pcc_message}"
         assert features_pass, f"Features comparison failed: {features_pcc_message}"
-
-    finally:
-        ttnn.close_device(device)