[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer

.github/workflows/pytorchsim_test.yml

@@ -166,6 +166,24 @@ jobs:
             -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \
             ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/conv/test_conv2d.py
 
+  test_conv_view_input:
+    name: Run test_conv_view_input.py
+    runs-on: ubuntu-latest
+    steps:
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Run test_conv_view_input.py
+        run: |
+          echo "Running test_conv_view_input.py"
+          docker run --rm \
+            -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \
+            ${{ inputs.image_name }} python3 PyTorchSim/tests/ops/conv/test_conv_view_input.py
+
   test_cat:
     name: Run test_cat.py
     runs-on: ubuntu-latest

PyTorchSimFrontend/mlir/mlir_conv_common.py

@@ -92,6 +92,19 @@ def outer_func_render(self, kernel_name, input_args):
         Y = self.output_node
         Bias = None if len(self.input_nodes) == 2 else self.input_nodes[2]
 
+        # The wrapper must not infer geometry from the runtime tensor: an input node can be
+        # a ReinterpretView, in which case call_kernel hands the wrapper the *base buffer*
+        # (possibly a different rank).  Bake the layout codegen assumed into the wrapper and
+        # rebuild the logical view from it, so the wrapper is correct for either argument.
+        def _layout(node):
+            return (tuple(int(s) for s in node.get_size()),
+                    tuple(int(s) for s in node.get_stride()),
+                    int(node.layout.offset))
+
+        X_size, X_stride, X_offset = _layout(X)
+        W_size, W_stride, W_offset = _layout(W)
+        _, _, I_H, I_W = X_size
+
         options = dict(
             kernel=self.kernel,
             KERNEL_NAME=kernel_name,
@@ -102,6 +115,16 @@ def outer_func_render(self, kernel_name, input_args):
             OUTPUT=Y,
             PADDING_H=self.padding[0],
             PADDING_W=self.padding[1],
+            X_SIZE=X_size,
+            X_STRIDE=X_stride,
+            X_OFFSET=X_offset,
+            W_SIZE=W_size,
+            W_STRIDE=W_stride,
+            W_OFFSET=W_offset,
+            I_H=I_H,
+            I_W=I_W,
+            X_PADDED_SIZE=(X_size[0], X_size[1],
+                           I_H + 2 * self.padding[0], I_W + 2 * self.padding[1]),
             VALIDATION_MODE=extension_config.pytorchsim_functional_mode,
             input_reorder=self.input_reorder
         )

PyTorchSimFrontend/mlir/mlir_conv_mt_template.py

@@ -100,12 +100,15 @@
 class MLIRConvMultiTileTemplate(MLIRConvCommonTemplate):
     WRAPPER_TEMPLATE = r"""
 def {{ FUNC_NAME }}{{kernel.def_wrapper()}}:
+    # Rebuild the logical NCHW inputs from the layout codegen assumed. The runtime argument
+    # may be the base buffer of a ReinterpretView (different rank/shape), so geometry is
+    # never inferred from X.shape.
+    X = reinterpret_tensor(X, {{ X_SIZE }}, {{ X_STRIDE }}, {{ X_OFFSET }})
+    W = reinterpret_tensor(W, {{ W_SIZE }}, {{ W_STRIDE }}, {{ W_OFFSET }})
+
     # Padding input
-    padded_shape = list(X.shape)
-    padded_shape[2] += 2 * {{ PADDING_H }}
-    padded_shape[3] += 2 * {{ PADDING_W }}
-    X_padding = torch.zeros(padded_shape).to(device=X.device)
-    X_padding[:, :, {{ PADDING_H }}:X.shape[2] + {{ PADDING_H }}, {{ PADDING_W }}:X.shape[3] + {{ PADDING_W }}] = X
+    X_padding = torch.zeros({{ X_PADDED_SIZE }}, device=X.device, dtype=X.dtype)
+    X_padding[:, :, {{ PADDING_H }}:{{ PADDING_H + I_H }}, {{ PADDING_W }}:{{ PADDING_W + I_W }}] = X
 
     # Tanspose inputs
     {%- for buf, name in kernel.get_conv_inputs().items() %}

PyTorchSimFrontend/mlir/mlir_conv_sb_template.py

@@ -101,12 +101,15 @@
 class MLIRConvSingleBatchTemplate(MLIRConvCommonTemplate):
     WRAPPER_TEMPLATE = r"""
 def {{ FUNC_NAME }}{{kernel.def_wrapper()}}:
+    # Rebuild the logical NCHW inputs from the layout codegen assumed. The runtime argument
+    # may be the base buffer of a ReinterpretView (different rank/shape), so geometry is
+    # never inferred from X.shape.
+    X = reinterpret_tensor(X, {{ X_SIZE }}, {{ X_STRIDE }}, {{ X_OFFSET }})
+    W = reinterpret_tensor(W, {{ W_SIZE }}, {{ W_STRIDE }}, {{ W_OFFSET }})
+
     # Padding input
-    padded_shape = list(X.shape)
-    padded_shape[2] += 2 * {{ PADDING_H }}
-    padded_shape[3] += 2 * {{ PADDING_W }}
-    X_padding = torch.zeros(padded_shape).to(device=X.device)
-    X_padding[:, :, {{ PADDING_H }}:X.shape[2] + {{ PADDING_H }}, {{ PADDING_W }}:X.shape[3] + {{ PADDING_W }}] = X
+    X_padding = torch.zeros({{ X_PADDED_SIZE }}, device=X.device, dtype=X.dtype)
+    X_padding[:, :, {{ PADDING_H }}:{{ PADDING_H + I_H }}, {{ PADDING_W }}:{{ PADDING_W + I_W }}] = X
 
     # Tanspose inputs
     {%- for buf, name in kernel.get_conv_inputs().items() %}

PyTorchSimFrontend/mlir/mlir_conv_sbs_template.py

@@ -101,12 +101,15 @@
 class MLIRConvSingleBatchStridedTemplate(MLIRConvCommonTemplate):
     WRAPPER_TEMPLATE = r"""
 def {{ FUNC_NAME }}{{kernel.def_wrapper()}}:
+    # Rebuild the logical NCHW inputs from the layout codegen assumed. The runtime argument
+    # may be the base buffer of a ReinterpretView (different rank/shape), so geometry is
+    # never inferred from X.shape.
+    X = reinterpret_tensor(X, {{ X_SIZE }}, {{ X_STRIDE }}, {{ X_OFFSET }})
+    W = reinterpret_tensor(W, {{ W_SIZE }}, {{ W_STRIDE }}, {{ W_OFFSET }})
+
     # Padding input
-    padded_shape = list(X.shape)
-    padded_shape[2] += 2 * {{ PADDING_H }}
-    padded_shape[3] += 2 * {{ PADDING_W }}
-    X_padding = torch.zeros(padded_shape).to(device=X.device)
-    X_padding[:, :, {{ PADDING_H }}:X.shape[2] + {{ PADDING_H }}, {{ PADDING_W }}:X.shape[3] + {{ PADDING_W }}] = X
+    X_padding = torch.zeros({{ X_PADDED_SIZE }}, device=X.device, dtype=X.dtype)
+    X_padding[:, :, {{ PADDING_H }}:{{ PADDING_H + I_H }}, {{ PADDING_W }}:{{ PADDING_W + I_W }}] = X
 
     # Tanspose inputs
     {%- for buf, name in kernel.get_conv_inputs().items() %}

PyTorchSimFrontend/mlir/mlir_conv_template.py

@@ -105,12 +105,15 @@
 class MLIRConvTemplate(MLIRConvCommonTemplate):
     WRAPPER_TEMPLATE = r"""
 def {{ FUNC_NAME }}{{kernel.def_wrapper()}}:
+    # Rebuild the logical NCHW inputs from the layout codegen assumed. The runtime argument
+    # may be the base buffer of a ReinterpretView (different rank/shape), so geometry is
+    # never inferred from X.shape.
+    X = reinterpret_tensor(X, {{ X_SIZE }}, {{ X_STRIDE }}, {{ X_OFFSET }})
+    W = reinterpret_tensor(W, {{ W_SIZE }}, {{ W_STRIDE }}, {{ W_OFFSET }})
+
     # Padding input
-    padded_shape = list(X.shape)
-    padded_shape[2] += 2 * {{ PADDING_H }}
-    padded_shape[3] += 2 * {{ PADDING_W }}
-    X_padding = torch.zeros(padded_shape).to(device=X.device)
-    X_padding[:, :, {{ PADDING_H }}:X.shape[2] + {{ PADDING_H }}, {{ PADDING_W }}:X.shape[3] + {{ PADDING_W }}] = X
+    X_padding = torch.zeros({{ X_PADDED_SIZE }}, device=X.device, dtype=X.dtype)
+    X_padding[:, :, {{ PADDING_H }}:{{ PADDING_H + I_H }}, {{ PADDING_W }}:{{ PADDING_W + I_W }}] = X
 
     # Tanspose inputs
     {%- for buf, name in kernel.get_conv_inputs().items() %}

tests/ops/conv/test_conv_view_input.py

@@ -0,0 +1,36 @@
+import os
+import sys
+import torch
+import torch._dynamo
+sys.path.insert(0, os.path.join(os.environ.get("TORCHSIM_DIR", default="/workspace/PyTorchSim"), "tests"))
+from _pytorchsim_utils import test_result
+
+def test_conv_view_input(device, batch_size=1, channels=32, height=8, width=8, out_channels=16, kernel_size=2):
+    """Conv whose input is a free view (ReinterpretView) over a lower-rank buffer.
+
+    A contiguous (B, H*W, C) buffer already is the channels_last layout of (B, C, H, W),
+    so permute+reshape is a pure view: inductor inserts no materializing copy and the conv
+    wrapper is handed the raw 3D base buffer. The wrapper must rebuild the 4D logical input
+    from the layout codegen assumed instead of reading X.shape. This is SegFormer's
+    efficient-attention spatial-reduction conv.
+    """
+    def conv_on_view(a, b, bias):
+        x = a * 2.0
+        h = x.permute(0, 2, 1).reshape(batch_size, channels, height, width)
+        return torch.nn.functional.conv2d(h, b, bias, stride=kernel_size, padding=0)
+    torch.manual_seed(0)
+    conv_input = torch.randn(batch_size, height * width, channels)
+    conv_kernel = torch.randn(out_channels, channels, kernel_size, kernel_size)
+    conv_bias = torch.randn(out_channels)
+    out = conv_on_view(conv_input, conv_kernel, conv_bias)
+    opt_fn = torch.compile(dynamic=False)(conv_on_view)
+    res = opt_fn(conv_input.to(device=device), conv_kernel.to(device=device), conv_bias.to(device=device))
+    test_result("Conv2d view input", res, out, rtol=1e-3, atol=1e-3)
+    print("Max diff > ", torch.max(torch.abs(res.cpu() - out)))
+
+if __name__ == "__main__":
+    device = torch.device("npu:0")
+    torch._dynamo.config.cache_size_limit = 64
+    with torch.no_grad():
+        test_conv_view_input(device)
+        test_conv_view_input(device, batch_size=1, channels=64, height=16, width=16, out_channels=32, kernel_size=4)