Refactors patch_merging to take weights from test

Author: ign-amanks

models/experimental/SSR/tests/test_basic_block.py

@@ -24,13 +24,12 @@ def to_2tuple(x):
     return (x, x)
 
 
-def create_basic_layer_preprocessor(device):
+def create_basic_layer_preprocessor(device, dim):
     def custom_preprocessor(torch_model, name, ttnn_module_args):
         params = {"blocks": {}}
 
         # Process each transformer block
         for i, block in enumerate(torch_model.blocks):
-            # relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0)
             params["blocks"][i] = preprocess_model_parameters(
                 initialize_model=lambda: block,
                 custom_preprocessor=create_swin_transformer_block_preprocessor(device),
@@ -41,7 +40,7 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
         if torch_model.downsample is not None:
             params["downsample"] = preprocess_model_parameters(
                 initialize_model=lambda: torch_model.downsample,
-                custom_preprocessor=create_patch_merging_preprocessor(device),
+                custom_preprocessor=create_patch_merging_preprocessor(device, dim),
                 device=device,
             )
 
@@ -88,7 +87,7 @@ def test_basic_layer(device, batch_size, input_resolution, dim, depth, num_heads
     # Create ttnn model
     params = preprocess_model_parameters(
         initialize_model=lambda: ref_layer,
-        custom_preprocessor=create_basic_layer_preprocessor(device),
+        custom_preprocessor=create_basic_layer_preprocessor(device, dim),
         device=device,
     )
 

models/experimental/SSR/tests/test_patch_merging.py

@@ -12,10 +12,31 @@
 from models.experimental.SSR.tt.patch_merging import TTPatchMerging
 
 
-def create_patch_merging_preprocessor(device):
+def create_patch_merging_preprocessor(device, dim):
     def custom_preprocessor(torch_model, name, ttnn_module_args):
         params = {}
 
+        # Create conv kernels for patch merging (same as in forward pass)
+        kernel_top_left = torch.zeros(dim, 1, 2, 2, dtype=torch.bfloat16)
+        kernel_top_left[:, 0, 0, 0] = 1.0
+
+        kernel_bottom_left = torch.zeros(dim, 1, 2, 2, dtype=torch.bfloat16)
+        kernel_bottom_left[:, 0, 1, 0] = 1.0
+
+        kernel_top_right = torch.zeros(dim, 1, 2, 2, dtype=torch.bfloat16)
+        kernel_top_right[:, 0, 0, 1] = 1.0
+
+        kernel_bottom_right = torch.zeros(dim, 1, 2, 2, dtype=torch.bfloat16)
+        kernel_bottom_right[:, 0, 1, 1] = 1.0
+
+        # Convert to TTNN tensors
+        params["conv_kernels"] = {
+            "top_left": ttnn.from_torch(kernel_top_left, device=device),
+            "bottom_left": ttnn.from_torch(kernel_bottom_left, device=device),
+            "top_right": ttnn.from_torch(kernel_top_right, device=device),
+            "bottom_right": ttnn.from_torch(kernel_bottom_right, device=device),
+        }
+
         # Linear reduction layer
         params["reduction"] = {
             "weight": ttnn.from_torch(
@@ -76,7 +97,7 @@ def test_patch_merging(device, batch_size, input_resolution, dim):
     # Create ttnn model
     params = preprocess_model_parameters(
         initialize_model=lambda: ref_layer,
-        custom_preprocessor=create_patch_merging_preprocessor(device),
+        custom_preprocessor=create_patch_merging_preprocessor(device, dim),
         device=device,
     )
 

models/experimental/SSR/tests/test_tile_selection.py

@@ -13,7 +13,7 @@
 from models.utility_functions import tt2torch_tensor, comp_pcc
 
 
-def create_tile_selection_preprocessor(device):
+def create_tile_selection_preprocessor(device, dim=96):
     def custom_preprocessor(torch_model, name, ttnn_module_args):
         parameters = {}
 
@@ -36,9 +36,10 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
         # Handle encoder layers - delegate to existing TTBasicLayer preprocessor
         if hasattr(torch_model, "layers"):
             for i, layer in enumerate(torch_model.layers):
+                layer_dim = int(dim * 2**i)
                 layer_params = preprocess_model_parameters(
                     initialize_model=lambda l=layer: l,
-                    custom_preprocessor=create_basic_layer_preprocessor(device),
+                    custom_preprocessor=create_basic_layer_preprocessor(device, layer_dim),
                     device=device,
                 )
                 parameters[f"layers.{i}"] = layer_params
@@ -127,16 +128,13 @@ def __init__(self, imgsz, patchsz, token_size, dim):
 
     parameters = preprocess_model_parameters(
         initialize_model=lambda: ref_layer,
-        custom_preprocessor=create_tile_selection_preprocessor(device),
+        custom_preprocessor=create_tile_selection_preprocessor(device, dim),
         device=device,
     )
 
     # Create TTNN implementation
     tt_layer = TTTileSelection(device=device, parameters=parameters, args=args, num_cls=num_cls)
 
-    # NCHW -> NHWC
-    input_tensor = input_tensor.permute(0, 2, 3, 1)
-
     # Convert input to TTNN
     tt_input = ttnn.from_torch(input_tensor, device=device, layout=ttnn.TILE_LAYOUT, dtype=ttnn.bfloat16)
 

models/experimental/SSR/tt/patch_merging.py

@@ -3,7 +3,6 @@
 
 import ttnn
 from models.common.lightweightmodule import LightweightModule
-import torch
 from models.demos.deepseek_v3.utils.config_helpers import matmul_config
 
 
@@ -27,6 +26,11 @@ def __init__(
         self.norm_weight = parameters["norm"]["weight"]
         self.norm_bias = parameters["norm"]["bias"]
 
+        self.kernel_top_left = parameters["conv_kernels"]["top_left"]
+        self.kernel_bottom_left = parameters["conv_kernels"]["bottom_left"]
+        self.kernel_top_right = parameters["conv_kernels"]["top_right"]
+        self.kernel_bottom_right = parameters["conv_kernels"]["bottom_right"]
+
     def forward(self, input_tensor):
         """
         Args:
@@ -44,96 +48,29 @@ def forward(self, input_tensor):
         input_tensor = ttnn.reshape(input_tensor, (B, H, W, C))
         x = ttnn.to_layout(input_tensor, ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.L1_MEMORY_CONFIG)
 
-        kernel_top_left = torch.zeros(C, 1, 2, 2, dtype=torch.bfloat16)
-        kernel_top_left[:, 0, 0, 0] = 1.0
-
-        kernel_bottom_left = torch.zeros(C, 1, 2, 2, dtype=torch.bfloat16)
-        kernel_bottom_left[:, 0, 1, 0] = 1.0
-
-        kernel_top_right = torch.zeros(C, 1, 2, 2, dtype=torch.bfloat16)
-        kernel_top_right[:, 0, 0, 1] = 1.0
-
-        kernel_bottom_right = torch.zeros(C, 1, 2, 2, dtype=torch.bfloat16)
-        kernel_bottom_right[:, 0, 1, 1] = 1.0
-
-        # Convert to TTNN tensors
-        tt_kernel_top_left = ttnn.from_torch(kernel_top_left, device=self.device)
-        tt_kernel_bottom_left = ttnn.from_torch(kernel_bottom_left, device=self.device)
-        tt_kernel_top_right = ttnn.from_torch(kernel_top_right, device=self.device)
-        tt_kernel_bottom_right = ttnn.from_torch(kernel_bottom_right, device=self.device)
-
-        # Apply grouped convolutions for each patch
-        x0 = ttnn.conv2d(
-            input_tensor=x,
-            weight_tensor=tt_kernel_top_left,
-            in_channels=C,
-            out_channels=C,
-            device=self.device,
-            kernel_size=(2, 2),
-            stride=(2, 2),
-            padding=(0, 0),
-            groups=C,  # Grouped convolution
-            batch_size=B,
-            input_height=H,
-            input_width=W,
-            conv_config=None,
-            dtype=ttnn.bfloat16,
-            memory_config=ttnn.DRAM_MEMORY_CONFIG,
-        )
-
-        x1 = ttnn.conv2d(
-            input_tensor=x,
-            weight_tensor=tt_kernel_bottom_left,
-            in_channels=C,
-            out_channels=C,
-            device=self.device,
-            kernel_size=(2, 2),
-            stride=(2, 2),
-            padding=(0, 0),
-            groups=C,
-            batch_size=B,
-            input_height=H,
-            input_width=W,
-            conv_config=None,
-            dtype=ttnn.bfloat16,
-            memory_config=ttnn.DRAM_MEMORY_CONFIG,
-        )
-
-        x2 = ttnn.conv2d(
-            input_tensor=x,
-            weight_tensor=tt_kernel_top_right,
-            in_channels=C,
-            out_channels=C,
-            device=self.device,
-            kernel_size=(2, 2),
-            stride=(2, 2),
-            padding=(0, 0),
-            groups=C,
-            batch_size=B,
-            input_height=H,
-            input_width=W,
-            conv_config=None,
-            dtype=ttnn.bfloat16,
-            memory_config=ttnn.DRAM_MEMORY_CONFIG,
-        )
-
-        x3 = ttnn.conv2d(
-            input_tensor=x,
-            weight_tensor=tt_kernel_bottom_right,
-            in_channels=C,
-            out_channels=C,
-            device=self.device,
-            kernel_size=(2, 2),
-            stride=(2, 2),
-            padding=(0, 0),
-            groups=C,
-            batch_size=B,
-            input_height=H,
-            input_width=W,
-            conv_config=None,
-            dtype=ttnn.bfloat16,
-            memory_config=ttnn.DRAM_MEMORY_CONFIG,
-        )
+        # Common convolution parameters
+        conv_params = {
+            "input_tensor": x,
+            "in_channels": C,
+            "out_channels": C,
+            "device": self.device,
+            "kernel_size": (2, 2),
+            "stride": (2, 2),
+            "padding": (0, 0),
+            "groups": C,  # Grouped convolution
+            "batch_size": B,
+            "input_height": H,
+            "input_width": W,
+            "conv_config": None,
+            "dtype": ttnn.bfloat16,
+            "memory_config": ttnn.DRAM_MEMORY_CONFIG,
+        }
+
+        # Apply grouped convolutions for each patch, this is instead of a slice operation
+        x0 = ttnn.conv2d(weight_tensor=self.kernel_top_left, **conv_params)
+        x1 = ttnn.conv2d(weight_tensor=self.kernel_bottom_left, **conv_params)
+        x2 = ttnn.conv2d(weight_tensor=self.kernel_top_right, **conv_params)
+        x3 = ttnn.conv2d(weight_tensor=self.kernel_bottom_right, **conv_params)
 
         ttnn.deallocate(x)