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())