Fix grouped conv1d conversion after decomposition (pytorch#22213)

Allow callers to configure graph-local weight producers that are safe
to convert. Configure Arm to accept the weight slices created by
grouped convolution decomposition.

Fix test_conv1d_groups for the arm_tosa_fp and arm_tosa_int flows. Add
TOSA FP and INT coverage for per-tensor and per-channel quantization.


Signed-off-by: Zingo Andersen <Zingo.Andersen@arm.com>

Author: zingo

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -16,6 +16,7 @@
 from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import (
     ConvertConv1dToConv2dPass,
 )
+from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -29,3 +30,12 @@ class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass):
         SizeAdjustInputPass,
     }
     target_ops = (exir_ops.edge.aten.convolution.default,)
+
+    def __init__(self, exported_program: ExportedProgram) -> None:
+        # Grouped-convolution decomposition creates one graph-local weight
+        # slice per group. Allow the shared pass to add the unit-height
+        # dimension after these producers.
+        super().__init__(
+            exported_program,
+            graph_local_weight_targets={exir_ops.edge.aten.slice_copy.Tensor},
+        )

backends/arm/test/ops/test_conv1d.py

@@ -217,6 +217,20 @@ def forward(self, x):
     batches=1,
 )
 
+# This is grouped but not depthwise: each group consumes two input channels.
+# Arm therefore decomposes it into three ordinary convolutions whose weights
+# are graph-local slices rather than lifted model parameters.
+conv1d_6_6x3x50_groups3 = Conv1d(
+    in_channels=6,
+    out_channels=6,
+    kernel_size=3,
+    stride=1,
+    padding=0,
+    groups=3,
+    length=50,
+    batches=4,
+)
+
 conv1d_7_1x3x16_st2_pd1_dl2 = Conv1d(
     in_channels=3,
     out_channels=3,
@@ -318,6 +332,42 @@ def test_convolution_1d_tosa_FP(test_data):
     pipeline.run()
 
 
+def test_convolution_1d_grouped_tosa_FP():
+    """Ensure sliced grouped-Conv1d weights are converted in the FP flow."""
+    pipeline = TosaPipelineFP[input_t](
+        conv1d_6_6x3x50_groups3,
+        conv1d_6_6x3x50_groups3.get_inputs(),
+        aten_op,
+        exir_op,
+    )
+    pipeline.run()
+
+
+@common.parametrize(
+    "per_channel_quantization", {"per_channel": True, "per_tensor": False}
+)
+def test_convolution_1d_grouped_tosa_INT(per_channel_quantization):
+    """Ensure sliced grouped-Conv1d weights are converted in the INT flow.
+
+    The decomposition must preserve both per-channel and per-tensor weight
+    quantization while converting each three-dimensional slice for Conv2d.
+
+    Args:
+        per_channel_quantization (bool): Whether weights use per-channel
+            instead of per-tensor quantization.
+
+    """
+    pipeline = TosaPipelineINT[input_t](
+        conv1d_6_6x3x50_groups3,
+        conv1d_6_6x3x50_groups3.get_inputs(),
+        aten_op,
+        exir_op,
+        per_channel_quantization=per_channel_quantization,
+        qtol=1,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_INT)
 def test_convolution_1d_tosa_INT(test_data):
     model, per_channel_quantization = test_data()

backends/transforms/convert_conv1d_to_conv2d_pass.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from collections.abc import Collection
+
 import torch
 
 from executorch.backends.transforms.utils import (
@@ -25,13 +27,30 @@ class ConvertConv1dToConv2dPass(ExportPass):
     are left unchanged.
     Weights passed across control-flow boundaries are unsqueezed inside the
     nested graph so its input signature remains unchanged.
+    Root graph-local weights are skipped by default. Backends may provide an
+    allow-list of safe producer targets when constructing the pass.
     The pass emits squeeze and unsqueeze boundaries for each converted
     convolution; downstream view cleanup can fold adjacent boundaries.
     """
 
-    def __init__(self, exported_program: ExportedProgram) -> None:
+    def __init__(
+        self,
+        exported_program: ExportedProgram,
+        graph_local_weight_targets: Collection[torch.fx.node.Target] = (),
+    ) -> None:
+        """Initialize the transform.
+
+        Args:
+            exported_program (ExportedProgram): Program containing the graph
+                and its lifted weights.
+            graph_local_weight_targets (Collection[torch.fx.node.Target]):
+                Producer targets whose three-dimensional outputs may be used
+                as Conv1d weights. All graph-local weights are skipped when
+                this is empty.
+        """
         super().__init__()
         self.exported_program = exported_program
+        self._graph_local_weight_targets = frozenset(graph_local_weight_targets)
 
     def _unsqueeze_weight(self, weight_node: torch.fx.Node) -> bool:
         weight = get_param_tensor(self.exported_program, weight_node)
@@ -182,6 +201,59 @@ def _convert_nested_graph(self, graph_module: torch.fx.GraphModule) -> bool:
             graph_module.recompile()
         return modified
 
+    def _convert_graph_local_weights(self, graph: torch.fx.Graph) -> bool:
+        """Convert supported root Conv1d nodes with graph-local weights.
+
+        A graph-local weight is produced by another graph operation rather
+        than stored as a model parameter. Callers choose which producers are
+        safe to convert. The original weight remains unchanged; an
+        unsqueeze is inserted between the producer and the convolution.
+
+        Args:
+            graph (torch.fx.Graph): Root graph whose Conv1d nodes are examined.
+
+        Returns:
+            bool: True if at least one Conv1d node was converted.
+        """
+        modified = False
+        for node in list(graph.nodes):
+            if not self._is_conv1d(node):
+                continue
+
+            input_node, weight_node = node.args[:2]
+            if not isinstance(input_node, torch.fx.Node) or not isinstance(
+                weight_node, torch.fx.Node
+            ):
+                continue
+
+            weight_meta = weight_node.meta.get("val")
+            weight_source_is_allowed = (
+                weight_node.target in self._graph_local_weight_targets
+            )
+            has_conv1d_weight_shape = (
+                isinstance(weight_meta, torch.Tensor) and weight_meta.dim() == 3
+            )
+            if not weight_source_is_allowed or not has_conv1d_weight_shape:
+                continue
+
+            # The weight producer must keep its original output shape. Insert
+            # the missing unit-height dimension only on the path to this conv:
+            #
+            #   graph-local weight [O, I, K]
+            #                 |
+            #             unsqueeze(2)
+            #                 v
+            #        Conv2d weight [O, I, 1, K]
+            self._convert_node(
+                graph,
+                node,
+                input_node,
+                weight_node,
+                unsqueeze_weight=True,
+            )
+            modified = True
+        return modified
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         is_root = graph_module is self.exported_program.graph_module
@@ -229,6 +301,12 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             )
             modified = True
 
+        # Some backends create graph-local weights known to be safe before this
+        # pass. A caller may opt those weights into the same local unsqueeze
+        # used by nested graphs without enabling arbitrary dynamic weights.
+        if is_root:
+            modified = self._convert_graph_local_weights(graph) or modified
+
         # The normal entry point invokes this call method for the root graph.
         # Explicitly apply the rewrite to its child control-flow GraphModules.
         if is_root:

backends/transforms/test/test_convert_conv1d_to_conv2d_pass.py

@@ -306,6 +306,38 @@ def test_skip_unfolded_qdq_weight():
     assert not result.modified
 
 
+@pytest.mark.parametrize("allow_slice_weight", [False, True])
+def test_graph_local_weight_target_opt_in(allow_slice_weight: bool):
+    model = Conv1d().eval()
+    x = torch.randn(1, 2, 8)
+    edge = _edge(model, (x,))
+    [conv] = [
+        node
+        for node in edge.graph.nodes
+        if node.target == exir_ops.edge.aten.convolution.default
+    ]
+    weight = conv.args[1]
+    with edge.graph.inserting_before(conv):
+        sliced_weight = edge.graph.call_function(
+            exir_ops.edge.aten.slice_copy.Tensor,
+            args=(weight, 0, 0, weight.meta["val"].shape[0]),
+        )
+    sliced_weight.meta["val"] = weight.meta["val"]
+    conv.replace_input_with(weight, sliced_weight)
+
+    allowed_targets = (
+        {exir_ops.edge.aten.slice_copy.Tensor} if allow_slice_weight else set()
+    )
+    result = ConvertConv1dToConv2dPass(
+        edge, graph_local_weight_targets=allowed_targets
+    ).call(edge.graph_module)
+
+    assert result.modified is allow_slice_weight
+    if allow_slice_weight:
+        assert conv.args[1].target == exir_ops.edge.aten.unsqueeze_copy.default
+        torch.testing.assert_close(edge.module()(x), model(x))
+
+
 def test_preserve_convolution_metadata():
     edge = _edge(Conv1d(), (torch.randn(1, 2, 8),))
     [conv] = [