[ET-VK] Rewrite conv1d only when the weight and the batch are ours to use

Two ways the rewrite could reach a graph it does not own, both from review.

The weight is reshaped in place, so a weight with a second consumer is
left behind at rank 4 while that consumer still expects rank 3. The
comment claimed the rank check in _eligible() covered this; it does not.
It stops the second conv1d being rewritten, which is precisely the
problem: the conv keeps its 1-D arguments and loses its 1-D weight. Only
rewrite when this convolution is the weight's sole consumer.

A symbolic batch was compared against 1, which reads the value it was
traced with rather than the values the graph accepts. A dimension free to
be 4 at runtime could pass a test for 1 and reach a conv2d that rejects
batched input. Require a batch that is already a plain 1.

Tests move to test_vulkan_passes.py, where the graph can be inspected
without a device. The delegate test only compared outputs, which an
unrewritten conv1d satisfies just as well, so it could not have failed if
the pass stopped firing.

The shared-weight test fails against the previous pass. The symbolic-batch
test does not: export specializes a size-1 dimension, so a batch that is
both symbolic and hinted 1 is not constructible here, and a batch hinted
higher is rejected either way. It pins the guard rather than reproducing
the defect.

Author: msluszniak

backends/vulkan/_passes/conv1d_as_conv2d.py

@@ -72,12 +72,25 @@ def _eligible(self, node: torch.fx.Node) -> bool:
         # The weight has to be a constant so it can be reshaped in place.
         if not is_param_node(self._exported_program, weight_node):
             return False
+        # ...and it has to belong to this convolution alone. The reshape is a
+        # mutation of shared state: a second consumer would keep reading the
+        # weight it was built against and find it rank 4, whether that consumer
+        # is another conv1d this pass then declines to rewrite, or an op that
+        # has nothing to do with convolution.
+        if len(weight_node.users) != 1:
+            return False
         transposed, groups = node.args[6], node.args[8]
         if transposed or groups != 1:
             return False
         # conv2d does not support batched input at all, so a rewrite there
         # would turn a working conv1d into a throw.
-        if in_shape[0] != 1:
+        #
+        # A symbolic batch is rejected rather than compared. Comparing it would
+        # read the value it happened to be traced with, so a dimension free to
+        # be 4 at runtime would pass a test for 1 and reach a conv2d that
+        # cannot serve it. Only a batch that is already a plain 1 is provably 1
+        # for every input the graph accepts.
+        if not isinstance(in_shape[0], int) or in_shape[0] != 1:
             return False
         # Depthwise and pointwise 1-D convs keep their own shaders.
         if w_shape[0] < _IM2COL_MIN_C_OUT or w_shape[2] == 1:
@@ -98,8 +111,9 @@ def _rewrite(self, graph: torch.fx.Graph, node: torch.fx.Node) -> None:
         out_4d_shape = [out_shape[0], out_shape[1], 1, out_shape[2]]
 
         # The weight tensor is contiguous, so inserting a singleton dim is a
-        # pure metadata change. A weight shared by several convs is only
-        # reshaped once; the rank check in _eligible() skips it afterwards.
+        # pure metadata change. It is safe to do in place only because
+        # _eligible() has established this convolution is the weight's sole
+        # consumer.
         weight = get_param_tensor(self._exported_program, weight_node)
         assert weight is not None
         set_param_tensor(

backends/vulkan/test/test_vulkan_passes.py

@@ -1,8 +1,9 @@
 import unittest
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 
+from executorch.backends.vulkan._passes.conv1d_as_conv2d import Conv1dAsConv2dPass
 from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass
 
 from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge
@@ -86,7 +87,104 @@ def op_node_count(graph_module: torch.fx.GraphModule, canonical_op_name: str) ->
     return count
 
 
+def run_conv1d_as_conv2d(
+    model: torch.nn.Module,
+    sample_inputs: Tuple[torch.Tensor],
+    dynamic_shapes=None,
+) -> torch.fx.GraphModule:
+    """Exports `model` to edge and runs Conv1dAsConv2dPass over it."""
+    program = torch.export.export(
+        model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True
+    )
+    edge_program = to_edge(
+        program,
+        compile_config=EdgeCompileConfig(_skip_dim_order=False, _check_ir_validity=False),
+    )
+    exported = edge_program.exported_program()
+
+    conv_pass = Conv1dAsConv2dPass()
+    conv_pass._exported_program = exported
+    return conv_pass(exported.graph_module).graph_module
+
+
+def conv_input_ranks(graph_module: torch.fx.GraphModule) -> List[int]:
+    """The rank of the input to every convolution in the graph."""
+    ranks = []
+    for node in graph_module.graph.nodes:
+        if get_target_canonical_name(node) == "convolution.default":
+            ranks.append(len(node.args[0].meta["val"].shape))
+    return ranks
+
+
 class TestVulkanPasses(unittest.TestCase):
+    def test_conv1d_as_conv2d_rewrites_eligible_conv(self):
+        # Output equality cannot show this: an unrewritten conv1d computes the
+        # same numbers. What distinguishes the two is the rank the convolution
+        # is asked to work at, so that is what is asserted.
+        class Conv1dModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv1d(32, 128, kernel_size=3, padding=1)
+
+            def forward(self, x):
+                return self.conv(x)
+
+        gm = run_conv1d_as_conv2d(
+            Conv1dModule(), (torch.randn(1, 32, 64),)
+        )
+
+        self.assertEqual(conv_input_ranks(gm), [4])
+        # Two views: one lifting the input to 4-D, one lowering the output back.
+        self.assertEqual(op_node_count(gm, "view_copy.default"), 2)
+
+    def test_conv1d_as_conv2d_rejects_shared_weight(self):
+        # The rewrite reshapes the weight in place. A second consumer would go
+        # on reading a tensor that has silently become rank 4, so a shared
+        # weight has to be left alone.
+        class SharedWeightModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(128, 32, 3))
+
+            def forward(self, x, y):
+                return (
+                    torch.nn.functional.conv1d(x, self.weight, padding=1),
+                    torch.nn.functional.conv1d(y, self.weight, padding=1),
+                )
+
+        gm = run_conv1d_as_conv2d(
+            SharedWeightModule(), (torch.randn(1, 32, 64), torch.randn(1, 32, 64))
+        )
+
+        self.assertEqual(conv_input_ranks(gm), [3, 3])
+        self.assertEqual(op_node_count(gm, "view_copy.default"), 0)
+
+    def test_conv1d_as_conv2d_rejects_symbolic_batch(self):
+        # A batch free to vary at runtime is not a batch of 1, whatever value it
+        # was traced with, and conv2d refuses batched input. The guard is on the
+        # dimension being a plain integer rather than on its traced value, so
+        # that a comparison never decides this from the hint alone.
+        #
+        # Traced at 2 because export specializes a size-1 dimension: the batch
+        # has to be genuinely symbolic here for the graph to be the one the
+        # guard exists for.
+        class Conv1dModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv1d(32, 128, kernel_size=3, padding=1)
+
+            def forward(self, x):
+                return self.conv(x)
+
+        gm = run_conv1d_as_conv2d(
+            Conv1dModule(),
+            (torch.randn(2, 32, 64),),
+            dynamic_shapes={"x": {0: torch.export.Dim.AUTO}},
+        )
+
+        self.assertEqual(conv_input_ranks(gm), [3])
+        self.assertEqual(op_node_count(gm, "view_copy.default"), 0)
+
     def test_fuse_torchao_quantized_embedding(self):
         """A torchao-dialect 4-bit weight-only quantized embedding
         (torchao.dequantize_affine -> aten.embedding) should fuse into a single