diff --git a/backends/arm/_passes/conv1d_unsqueeze_pass.py b/backends/arm/_passes/conv1d_unsqueeze_pass.py index 591a14c455b..d432e5e37f9 100644 --- a/backends/arm/_passes/conv1d_unsqueeze_pass.py +++ b/backends/arm/_passes/conv1d_unsqueeze_pass.py @@ -15,6 +15,8 @@ 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 @@ -26,3 +28,12 @@ class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass): RewriteConvPass, SizeAdjustInputPass, } + + 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}, + ) diff --git a/backends/arm/test/ops/test_conv1d.py b/backends/arm/test/ops/test_conv1d.py index 5fcdcf50465..775a039d099 100644 --- a/backends/arm/test/ops/test_conv1d.py +++ b/backends/arm/test/ops/test_conv1d.py @@ -212,6 +212,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, @@ -313,6 +327,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() diff --git a/backends/transforms/convert_conv1d_to_conv2d_pass.py b/backends/transforms/convert_conv1d_to_conv2d_pass.py index d297f771155..de4e6860fea 100644 --- a/backends/transforms/convert_conv1d_to_conv2d_pass.py +++ b/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: diff --git a/backends/transforms/test/test_convert_conv1d_to_conv2d_pass.py b/backends/transforms/test/test_convert_conv1d_to_conv2d_pass.py index 56ef534a402..3de2d807529 100644 --- a/backends/transforms/test/test_convert_conv1d_to_conv2d_pass.py +++ b/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] = [