diff --git a/backends/vulkan/_passes/__init__.py b/backends/vulkan/_passes/__init__.py index 1afaf48dde7..1369610ae9f 100644 --- a/backends/vulkan/_passes/__init__.py +++ b/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", diff --git a/backends/vulkan/_passes/normalize_convolution_args.py b/backends/vulkan/_passes/normalize_convolution_args.py new file mode 100644 index 00000000000..37bda323ed4 --- /dev/null +++ b/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) diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index f030b9268a1..0c873b4f0ad 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/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]) diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index f7d6955ce26..de2b1528fa5 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/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(),