Normalize single element convolution arg lists

backends/vulkan/_passes/__init__.py

@@ -15,6 +15,9 @@
     InsertDtypePromotionPass,
 )
 from executorch.backends.vulkan._passes.insert_prepack_nodes import insert_prepack_nodes
+from executorch.backends.vulkan._passes.normalize_convolution_args import (
+    NormalizeConvolutionArgs,
+)
 from executorch.backends.vulkan._passes.remove_asserts import (
     remove_asserts,
     RemoveAssertsTransform,
@@ -33,6 +36,7 @@
     "FuseQuantizedOpsTransform",
     "InsertDtypePromotionPass",
     "insert_prepack_nodes",
+    "NormalizeConvolutionArgs",
     "remove_asserts",
     "RemoveAssertsTransform",
     "RemoveRedundantOpsTransform",

backends/vulkan/_passes/normalize_convolution_args.py

@@ -0,0 +1,73 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from typing import Dict, List, Tuple
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
+
+from torch.fx.node import Argument
+
+
+class NormalizeConvolutionArgs(ExportPass):
+    """
+    Broadcasts single element stride/padding/dilation/output_padding lists of a
+    convolution to the number of spatial dimensions.
+
+    ATen allows these arguments to be given as a single value that applies to
+    every spatial dimension, and torch.nn does exactly that for
+    `padding="valid"`, which exports as `padding=[0]` rather than `[0, 0]`.
+    The Vulkan convolution reads them as fixed width vectors
+    (`make_ivec2_from_list` -> `make_ivec2`, which requires exactly 2 elements),
+    so a 2D convolution written that way aborts at the first inference:
+
+      make_ivec2 ... (ints.size() == 2) is false!
+
+    Normalizing here keeps the graph ATen compliant and leaves the runtime
+    unchanged.
+    """
+
+    # arg index -> name, for the list arguments of aten.convolution
+    _list_arg_indices: Dict[int, str] = {
+        3: "stride",
+        4: "padding",
+        5: "dilation",
+        7: "output_padding",
+    }
+
+    def call_operator(
+        self,
+        op,  # pyre-ignore
+        args: Tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        if op != exir_ops.edge.aten.convolution.default:
+            return super().call_operator(op, args, kwargs, meta)
+
+        # weight is (out_channels, in_channels / groups, *kernel_size)
+        weight = args[1]
+        # pyre-ignore[16]
+        spatial_dims = len(weight.node.meta["val"].shape) - 2
+        if spatial_dims < 2:
+            return super().call_operator(op, args, kwargs, meta)
+
+        new_args: List[Argument] = list(args)
+        modified = False
+        for idx in self._list_arg_indices:
+            if idx >= len(new_args):
+                continue
+            value = new_args[idx]
+            if isinstance(value, (list, tuple)) and len(value) == 1:
+                new_args[idx] = [value[0]] * spatial_dims
+                modified = True
+
+        if not modified:
+            return super().call_operator(op, args, kwargs, meta)
+
+        return super().call_operator(op, tuple(new_args), kwargs, meta)

backends/vulkan/test/test_vulkan_passes.py

@@ -779,3 +779,67 @@ def forward(self, x):
 
         gm = ep.graph_module
         self.assertEqual(op_node_count(gm, "q8ta_pixel_shuffle.default"), 0)
+
+    def test_normalize_convolution_args_broadcasts_single_element_lists(self):
+        """A 2D convolution written `padding="valid"` exports a one element
+        padding list, which the Vulkan convolution cannot read.
+
+        make_ivec2_from_list -> make_ivec2 requires exactly 2 elements, so such
+        a convolution aborted at the first inference with
+
+          make_ivec2 ... (ints.size() == 2) is false!
+
+        This is how torch.nn writes `padding="valid"`, and how HuggingFace
+        writes SigLIP/CLIP patch embeddings, so every vision tower of that shape
+        was affected.
+        """
+        from executorch.backends.vulkan._passes.normalize_convolution_args import (
+            NormalizeConvolutionArgs,
+        )
+        from executorch.exir.program._program import _transform
+
+        class Conv2dValid(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 16, kernel_size=3, padding="valid")
+
+            def forward(self, x):
+                return self.conv(x)
+
+        class Conv1dValid(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv1d(3, 16, kernel_size=3, padding="valid")
+
+            def forward(self, x):
+                return self.conv(x)
+
+        def padding_of(program) -> list:
+            conv = next(
+                node
+                for node in program.graph_module.graph.nodes
+                if get_target_canonical_name(node) == "convolution.default"
+            )
+            return list(conv.args[4])
+
+        # 2D: one element padding must be broadcast to the two spatial dims.
+        edge_program = to_edge(
+            torch.export.export(Conv2dValid().eval(), (torch.randn(1, 3, 32, 32),)),
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+        ep = edge_program._edge_programs["forward"]
+        self.assertEqual(padding_of(ep), [0], "expected export to emit a 1D padding")
+
+        ep = _transform(ep, NormalizeConvolutionArgs())
+        self.assertEqual(padding_of(ep), [0, 0])
+
+        # 1D: a one element list already matches the single spatial dim, so it
+        # must be left alone.
+        edge_program = to_edge(
+            torch.export.export(Conv1dValid().eval(), (torch.randn(1, 3, 32),)),
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+        ep = _transform(
+            edge_program._edge_programs["forward"], NormalizeConvolutionArgs()
+        )
+        self.assertEqual(padding_of(ep), [0])

backends/vulkan/vulkan_preprocess.py

@@ -21,6 +21,7 @@
     FuseQuantizedOpsTransform,
     insert_prepack_nodes,
     InsertDtypePromotionPass,
+    NormalizeConvolutionArgs,
     RemoveRedundantOpsTransform,
     SqueezeUnsqueezeInputs,
     TagMemoryMetaPass,
@@ -185,6 +186,7 @@ def preprocess(  # noqa: C901
         program = apply_passes(
             program,
             [
+                NormalizeConvolutionArgs(),
                 AddmmToLinearTransform(),
                 FuseBatchNormPass(program),
                 AddmmToLinearTransform(),