Lower Conv1d atomically through TOSA Conv2d (pytorch#22282)

Differential Revision: D116686805

Pull Request resolved: pytorch#22282

Author: mcremon-meta

backends/arm/_passes/BUCK

@@ -67,7 +67,6 @@ fbcode_target(
         "//executorch/backends/transforms:fuse_duplicate_users_pass",
         "//executorch/backends/transforms:propagate_view_copy_permute_pass",
         "//executorch/backends/transforms:fuse_identical_input_transforms_pass",
-        "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass",
         "//executorch/backends/transforms:fuse_view_copy",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/transforms:replace_scalar_with_tensor",

backends/arm/_passes/__init__.py

@@ -20,7 +20,6 @@
 from .cast_int_comparison_inputs_pass import CastIntComparisonInputsPass  # noqa
 from .cast_to_int32_pass import CastToInt32Pass  # noqa
 from .constant_folding_pass import ConstantFoldingPass  # noqa
-from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass  # noqa
 from .convert_elu_params import ConvertELUParamsPass  # noqa
 from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass  # noqa
 from .convert_full_like_to_full_pass import ConvertFullLikeToFullPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -22,7 +22,6 @@
     ComputeConstantOpsAOTPass,
     ConstantFoldingPass,
     ControlFlowConstInlinePass,
-    Conv1dUnsqueezePass,
     ConvertEluFamilyToEluPass,
     ConvertELUParamsPass,
     ConvertExpandCopyToRepeatPass,
@@ -603,7 +602,6 @@ def _tosa_pipeline(
                 DecomposeAdaptiveAvgPool2dPass(),
                 DecomposeDynamicAdaptiveAvgPool2dPass(),
                 DecomposeAvgPool2dPass(),
-                Conv1dUnsqueezePass(exported_program),
             ]
         )
 

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -8,8 +8,8 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
-from executorch.backends.arm._passes.conv1d_unsqueeze_pass import Conv1dUnsqueezePass
 from executorch.backends.arm._passes.quant_args import QuantArgs
+from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -46,7 +46,7 @@ class DecomposeGroupedConvPass(ArmOpTargetedPass):
 
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = {Conv1dUnsqueezePass}
+    _passes_required_after: Set[Type[ExportPass]] = {RewriteConvPass}
     target_ops = (
         exir_ops.edge.aten.convolution.default,
         torch.ops.aten.conv_transpose2d.input,

backends/arm/_passes/decompose_grouped_conv_pass.py

@@ -94,7 +94,7 @@ def _adjust_pad_if_needed(
                 pass instead.
 
         """
-        mod_remainder = (
+        mod_remainder: int | torch.SymInt = (
             input_len + 2 * pad - dilation * (input_weight - 1) - 1
         ) % stride
 
@@ -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]
@@ -524,7 +524,7 @@ def _combine_rescale_scales(
     @staticmethod
     def _is_direct_int32_rescale(node: torch.fx.Node) -> bool:
         """Return whether a node directly rescales its input to INT32."""
-        return (
+        return bool(
             node.op == "call_function"
             and node.target == exir_ops.backend.tosa.RESCALE.default
             and len(node.args) > 1
@@ -911,7 +911,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 +997,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 +1102,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,
@@ -1059,16 +1143,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                 tosa_node_fake_tensor.dtype == torch.int32
                 and input_fake_tensor.dtype == torch.int16
             )
-            if is_a16w8_conv:
+            if is_a16w8_conv and spatial_rank != 1:
                 # Keep values in INT32 whenever a consumer supports it, even
                 # though the declared output is INT16, by branching from the
                 # accumulator before narrowing.
+                #
+                # Rank-three convolutions are excluded. The legacy Conv1d
+                # expansion placed a rank-changing view between the convolution
+                # and its INT32 consumers, so the convolution narrowed to its
+                # exported output domain instead of forking. Forking here would
+                # give each branch its own boundary rescale and permute, which
+                # Vela materialises as a second full transpose of the output.
                 self._insert_a16w8_output_branches(
                     graph_module,
                     node,
                     tosa_op,
                     tosa_node_fake_tensor,
-                    post_permute_input,
+                    output_conversion_node,
                     post_permute_dims,
                 )
                 # Only users not moved to widened branches remain on the
@@ -1078,7 +1169,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)
 

backends/arm/_passes/rewrite_conv_pass.py

@@ -298,6 +298,38 @@ def test_canonicalize_moves_permute_before_view() -> None:
     _validate_numerics(gm_before, result.graph_module, (x_data,))
 
 
+def test_canonicalize_sinks_singleton_view_below_permute() -> None:
+    builder = GraphBuilder()
+    x_data = torch.randn(2, 8, 64)
+    x = builder.placeholder("x", x_data)
+    v1 = builder.call_operator(
+        op=exir_ops.edge.aten.view_copy.default,
+        args=(x, [2, 8, 1, 64]),
+    )
+    p1 = builder.call_operator(
+        op=exir_ops.edge.aten.permute_copy.default,
+        args=(v1, [0, 2, 3, 1]),
+    )
+    builder.output([p1])
+    original = builder.get_graph_module()
+    gm_before = copy.deepcopy(original)
+
+    pass_instance = CanonicalizeViewCopyPermutePass()
+    result = cast(PassResult, pass_instance.call(original))
+
+    assert result.modified
+    compute_nodes = [
+        node for node in result.graph_module.graph.nodes if node.op == "call_function"
+    ]
+    assert [node.target for node in compute_nodes] == [
+        exir_ops.edge.aten.permute_copy.default,
+        exir_ops.edge.aten.view_copy.default,
+    ]
+    assert compute_nodes[0].args[1] == [0, 2, 1]
+    assert compute_nodes[1].args[1] == [2, 1, 64, 8]
+    _validate_numerics(gm_before, result.graph_module, (x_data,))
+
+
 def test_canonicalize_follows_interleaved_chain_users() -> None:
     builder = GraphBuilder()
     x_data = torch.randn(4, 2, 4)

backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.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")