From 8499c669cbf6d5e2bbe755d5a01ba5ba1597732a Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Sun, 30 Aug 2026 12:26:09 -0700 Subject: [PATCH] Lower Conv1d atomically through TOSA Conv2d (#22282) Summary: Lower one-dimensional convolutions atomically during convolution rewriting instead of expanding them to rank four in an earlier pass. Emit rank-three layout boundaries around the operation and introduce the singleton spatial dimension only after converting the input to the backend layout. Rewrite regular and depthwise weights directly into their TOSA Conv2d layouts, and retain the earlier expansion path only for transposed one-dimensional convolutions. This exposes rank-three transforms to the existing layout propagation and cancellation passes while keeping the generated convolution representation valid. Differential Revision: D116686805 --- backends/arm/_passes/conv1d_unsqueeze_pass.py | 19 ++- backends/arm/_passes/rewrite_conv_pass.py | 122 +++++++++++++++++- ...ve_permutes_around_elementwise_tosa_ops.py | 36 ++++++ .../arm/test/passes/test_rewrite_conv_pass.py | 120 +++++++++++++++++ 4 files changed, 289 insertions(+), 8 deletions(-) diff --git a/backends/arm/_passes/conv1d_unsqueeze_pass.py b/backends/arm/_passes/conv1d_unsqueeze_pass.py index 7bc86a4fcf2..44c968880fa 100644 --- a/backends/arm/_passes/conv1d_unsqueeze_pass.py +++ b/backends/arm/_passes/conv1d_unsqueeze_pass.py @@ -7,6 +7,8 @@ from typing import Set, Type +import torch + from executorch.backends.arm._passes import ArmOpTargetedPass from executorch.backends.arm._passes.convert_squeezes_to_view import ( ConvertSqueezesToViewPass, @@ -21,7 +23,14 @@ class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass): - """Arm wrapper for the shared Conv1d-to-Conv2d transform.""" + """Convert ConvTranspose1d into ConvTranspose2d. + + Forward Conv1d is lowered atomically by ``RewriteConvPass`` so its layout + transforms can be emitted in the rank-3 domain. TOSA has no native + transpose-convolution 1D operator, so this pass retains the rank-4 + expansion for ConvTranspose1d. + + """ _passes_required_after: Set[Type[ExportPass]] = { ConvertSqueezesToViewPass, @@ -29,3 +38,11 @@ class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass): SizeAdjustInputPass, } target_ops = (exir_ops.edge.aten.convolution.default,) + + def _conv1d_weight_node(self, node: torch.fx.Node) -> torch.fx.Node | None: + weight_node = super()._conv1d_weight_node(node) + # A non-None weight node means the base pass matched the full nine-arg + # convolution signature, so args[6] (transposed) is always present. + if weight_node is None or not node.args[6]: + return None + return weight_node diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index 69d4d76c9cf..5e6da5639f4 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -121,14 +121,14 @@ def _adjust_pad_if_needed( return pad - mod_remainder - def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool: + def _is_depthwise_conv(self, node: torch.fx.Node) -> bool: if ( node.op != "call_function" or node.target != exir_ops.edge.aten.convolution.default ): return False input_tensor = get_first_fake_tensor(node.all_input_nodes[0]) - if len(input_tensor.shape) != 4: + if len(input_tensor.shape) not in (3, 4): return False groups = node.args[-1] in_channels = input_tensor.shape[1] @@ -558,8 +558,26 @@ def _insert_layout_permute( input_node: torch.fx.Node, input_fake_tensor: FakeTensor, dims: tuple[int, ...], + squeeze_spatial: bool = False, ) -> tuple[torch.fx.Node, FakeTensor]: - """Insert the mandatory TOSA-to-Edge output layout permutation.""" + """Insert the mandatory TOSA-to-Edge output layout permutation. + + ``squeeze_spatial`` drops the singleton spatial dimension a rank-three + convolution carries through TOSA Conv2d, so the layout permutation is + emitted in the rank-three domain the exported graph expects. + + """ + if squeeze_spatial: + squeezed_fake_tensor = cast(FakeTensor, input_fake_tensor.squeeze(1)) + with graph_module.graph.inserting_after(input_node): + input_node = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(input_node, list(squeezed_fake_tensor.shape)), + from_node=source_node, + ) + input_node.meta["val"] = squeezed_fake_tensor + input_fake_tensor = squeezed_fake_tensor with graph_module.graph.inserting_after(input_node): output = create_node( graph=graph_module.graph, @@ -579,6 +597,7 @@ def _insert_a16w8_output_branches( tosa_node_fake_tensor: torch.Tensor, default_rescale: torch.fx.Node, post_permute_dims: tuple[int, ...], + squeeze_spatial: bool = False, ) -> None: """Route A16W8 convolution users through the required output types. @@ -652,6 +671,7 @@ def _insert_a16w8_output_branches( int32_user, int32_fake_tensor, post_permute_dims, + squeeze_spatial, ) for user in previous_users: user.replace_input_with(int32_user, int32_output) @@ -702,6 +722,7 @@ def _insert_a16w8_output_branches( widened_rescale, widened_fake_tensor, post_permute_dims, + squeeze_spatial, ) # Clone the source permutation for the widened branch. The @@ -911,7 +932,70 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 dilation = tuple(dilation_list) pad = pad_attr - if self._is_conv3d(len(input_shape), group): + if spatial_rank == 1: + target_op = ( + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default + if self._is_depthwise_conv(node) + else exir_ops.backend.tosa.CONV2D.default + ) + pre_permute_dims = (0, 2, 1) + post_permute_dims = (0, 2, 1) + with graph_module.graph.inserting_before(node): + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.permute_copy.default, + args=(x, list(pre_permute_dims)), + from_node=node, + ) + permuted_input_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims + ) + x.meta["val"] = permuted_input_fake + input_tensor_for_tosa_fake = permuted_input_fake.unsqueeze(1) + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(x, list(input_tensor_for_tosa_fake.shape)), + from_node=node, + ) + x.meta["val"] = input_tensor_for_tosa_fake + + kernel_width = weight_shape[2] + if target_op == exir_ops.backend.tosa.DEPTHWISE_CONV2D.default: + in_channels = input_fake_tensor.shape[1] + channel_multiplier = weight_shape[0] // in_channels + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(1, 2, 0), + name_suffix="hwicm", + reshape_dims=( + 1, + kernel_width, + in_channels, + channel_multiplier, + ), + ) + else: + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(0, 2, 1), + name_suffix="ohwi", + reshape_dims=( + weight_shape[0], + 1, + kernel_width, + weight_shape[1], + ), + ) + weight_fake_tensor = get_first_fake_tensor(weight) + stride = (1, stride[0]) + dilation = (1, dilation[0]) + pad = [0, 0, pad[0], pad[1]] + elif self._is_conv3d(len(input_shape), group): target_op = exir_ops.backend.tosa.CONV3D.default pre_permute_dims = ODHWI_ORDER post_permute_dims = ODHWI_INVERSE_ORDER @@ -934,7 +1018,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 name_suffix="odhwi", ) weight_fake_tensor = get_first_fake_tensor(weight) - elif self._is_depthwise_conv2d(node): + elif self._is_depthwise_conv(node): target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default pre_permute_dims = NHWC_ORDER post_permute_dims = NHWC_INVERSE_ORDER @@ -1039,7 +1123,28 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 if post_permute_dims is None: raise RuntimeError("Expected post permute dims for explicit layout") + output_conversion_node = node_replacement post_permute_input = node_replacement + squeeze_view: torch.fx.Node | None = None + if spatial_rank == 1: + squeezed_output_fake = cast( + FakeTensor, node_replacement_fake_tensor.squeeze(1) + ) + special_dtype = node_replacement.meta.get(TosaSpecialDtype.meta_key()) + with graph_module.graph.inserting_after(node_replacement): + node_replacement = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(node_replacement, list(squeezed_output_fake.shape)), + from_node=node, + ) + node_replacement.meta["val"] = squeezed_output_fake + if special_dtype: + node_replacement.meta[TosaSpecialDtype.meta_key()] = special_dtype + squeeze_view = node_replacement + post_permute_input = node_replacement + node_replacement_fake_tensor = squeezed_output_fake + with graph_module.graph.inserting_after(node_replacement): node_replacement = create_node( graph=graph_module.graph, @@ -1068,8 +1173,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 node, tosa_op, tosa_node_fake_tensor, - post_permute_input, + output_conversion_node, post_permute_dims, + spatial_rank == 1, ) # Only users not moved to widened branches remain on the # source node. Route those through the declared INT16 output, @@ -1078,7 +1184,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 node.replace_all_uses_with(node_replacement) else: graph_module.graph.erase_node(node_replacement) - graph_module.graph.erase_node(post_permute_input) + if squeeze_view is not None: + graph_module.graph.erase_node(squeeze_view) + graph_module.graph.erase_node(output_conversion_node) else: node.replace_all_uses_with(node_replacement) diff --git a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py index 864b6c669f9..f7af7cc41e2 100644 --- a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py @@ -26,6 +26,8 @@ RESCALE_TARGET = exir_ops.backend.tosa.RESCALE.default MUL_TARGET = exir_ops.edge.aten.mul.Tensor ADD_TARGET = exir_ops.edge.aten.add.Tensor +SUB_TARGET = exir_ops.edge.aten.sub.Tensor +VIEW_TARGET = exir_ops.edge.aten.view_copy.default ERF_TARGET = exir_ops.edge.aten.erf.default @@ -150,6 +152,40 @@ def test_remove_permutes_around_rescale_tosa_INT() -> None: assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1 +def test_sink_view_preserves_layout_through_rescale_to_broadcast_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 4, 1, 1) + direct = graph.placeholder("direct") + direct.meta["val"] = torch.randn(1, 8, 4) + + permute = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.randn(1, 1, 1, 4) + mul = graph.create_node("call_function", MUL_TARGET, args=(permute, permute)) + mul.meta["val"] = torch.randn(1, 1, 1, 4) + sink = graph.create_node("call_function", VIEW_TARGET, args=(mul, [1, 1, 4])) + sink.meta["val"] = torch.randn(1, 1, 4) + rescale = graph.create_node( + "call_function", + RESCALE_TARGET, + args=(sink, torch.int8, [1.0], 0, 0), + ) + rescale.meta["val"] = torch.randn(1, 1, 4) + sub = graph.create_node("call_function", SUB_TARGET, args=(direct, rescale)) + sub.meta["val"] = torch.randn(1, 8, 4) + graph.output(sub) + + graph_module = torch.fx.GraphModule({}, graph) + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert not result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 1 + assert sub.args == (direct, rescale) + + def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/arm/test/passes/test_rewrite_conv_pass.py b/backends/arm/test/passes/test_rewrite_conv_pass.py index 31c4205f16f..efd41db22ed 100644 --- a/backends/arm/test/passes/test_rewrite_conv_pass.py +++ b/backends/arm/test/passes/test_rewrite_conv_pass.py @@ -90,6 +90,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) + x +class A16W8Conv1dInt32Consumer(nn.Module): + """Exercise a rank-three A16W8 convolution consumed only by an INT32 add.""" + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Feed the convolution output directly to a residual addition.""" + return self.conv(x) + x + + class A16W8MixedConsumer(nn.Module): """Exercise a shared A16W8 convolution output with mixed consumers.""" @@ -296,6 +308,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dBiasModule(torch.nn.Module): + def __init__(self, depthwise: bool = False) -> None: + super().__init__() + groups = 4 if depthwise else 1 + out_channels = 8 if depthwise else 6 + self.conv = torch.nn.Conv1d( + 4, + out_channels, + kernel_size=3, + padding=1, + groups=groups, + bias=True, + ) + + def get_inputs(self) -> tuple[torch.Tensor]: + return (torch.randn(1, 4, 8),) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + class Conv3dBiasModule(torch.nn.Module): def __init__(self) -> None: super().__init__() @@ -491,6 +524,37 @@ def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: assert direct_int32_rescales[0].args[2] == pytest.approx(expected_int32_scales[0]) +def test_rewrite_conv1d_a16w8_squeezes_int32_branch_to_rank3() -> None: + """Test that the widened INT32 branch of a Conv1d lands back in rank three. + + The INT32 branch forks off the INT48 accumulator and needs its own layout + boundary. TOSA Conv2d produces a rank-four NHWC tensor for a rank-three + convolution, so the branch must drop the singleton spatial dimension before + the rank-three output permutation. + + """ + inputs = (torch.randn(1, 4, 8),) + gm, expected_int32_scales = _rewrite_a16w8_convs(A16W8Conv1dInt32Consumer(), inputs) + + direct_int32_rescales = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.backend.tosa.RESCALE.default + and node.args[1] == torch.int32 + and node.all_input_nodes[0].target == exir_ops.backend.tosa.CONV2D.default + ] + assert len(direct_int32_rescales) == len(expected_int32_scales) == 1 + assert direct_int32_rescales[0].args[2] == pytest.approx(expected_int32_scales[0]) + + (branch_view,) = tuple(direct_int32_rescales[0].users) + assert branch_view.target == exir_ops.edge.aten.view_copy.default + assert branch_view.meta["val"].shape == torch.Size((1, 8, 4)) + (branch_permute,) = tuple(branch_view.users) + assert branch_permute.target == exir_ops.edge.aten.permute_copy.default + assert branch_permute.meta["val"].shape == torch.Size((1, 4, 8)) + + def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: r"""Test that an indirect INT32 consumer keeps a widened branch. @@ -533,6 +597,62 @@ def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: assert len(widened_paths) == 1 +@pytest.mark.parametrize( + "depthwise,target_op,expected_weight_shape,expected_output_shape", + [ + ( + False, + exir_ops.backend.tosa.CONV2D.default, + (6, 1, 3, 4), + (1, 6, 8), + ), + ( + True, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + (1, 3, 4, 2), + (1, 8, 8), + ), + ], +) +def test_rewrite_conv1d_emits_atomic_rank3_layout_boundaries( + depthwise: bool, + target_op, + expected_weight_shape: tuple[int, ...], + expected_output_shape: tuple[int, ...], +) -> None: + module = Conv1dBiasModule(depthwise).eval() + edge_program = to_edge(export(module, module.get_inputs())).exported_program() + + with TosaLoweringContext(_compile_spec().tosa_spec): + result = RewriteConvPass(edge_program)(edge_program.graph_module) + assert result is not None + graph_module = result.graph_module + + conv = _get_call_function_node(graph_module, target_op) + input_view = conv.args[0] + assert isinstance(input_view, torch.fx.Node) + assert input_view.target == exir_ops.edge.aten.view_copy.default + input_permute = input_view.args[0] + assert isinstance(input_permute, torch.fx.Node) + assert input_permute.target == exir_ops.edge.aten.permute_copy.default + assert input_permute.args[1] == [0, 2, 1] + assert input_view.meta["val"].shape == torch.Size((1, 1, 8, 4)) + + weight = conv.args[1] + assert isinstance(weight, torch.fx.Node) + assert weight.meta["val"].shape == torch.Size(expected_weight_shape) + + output_view = next( + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.view_copy.default and node.args[0] is conv + ) + output_permute = next(iter(output_view.users)) + assert output_permute.target == exir_ops.edge.aten.permute_copy.default + assert output_permute.args[1] == [0, 2, 1] + assert output_permute.meta["val"].shape == torch.Size(expected_output_shape) + + @pytest.mark.skipif(not _VGF_ENABLED, reason="VGF not enabled") def test_fold_and_annotate_q_params_vgf_quant_tracks_fused_relu_qparams() -> None: exported_program = _export_quantized(TinyConvReluCat())