Arm backend: Fix pad transpose handling around grid sampler (pytorch#21738)

Arm backend: Fix pad transpose handling around grid sampler
Change-Id: I5e9282082c670a086262c13f68bdecbd44ec5a05
Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @rascani

Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>

Author: robell

backends/arm/_passes/arm_pass_manager.py

@@ -649,10 +649,10 @@ def _tosa_pipeline(
                 RewriteMXFPConv2dPass(exported_program),
                 RewriteMXFPLinearPass(exported_program),
                 RewriteMatmulPass(),
-                RewritePadPass(),
                 FuseViewCopyTransformPass(),
                 PropagateViewCopyPermuteDownPass(self.compile_spec, exported_program),
                 PropagateViewCopyPermuteUpPass(self.compile_spec, exported_program),
+                RewritePadPass(),
                 # Propagation can leave a binary op with mismatched operand ranks,
                 # which TOSA rejects; re-match ranks before lowering.
                 MoveDataMovementOpsToSmallerDtypePass(),

backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py

@@ -420,7 +420,10 @@ def test_rewrite_grid_sampler_to_tosa_custom_c3_pad_for_align_corners():
     )
     assert payload["output_0_type"] == "Image"
     assert payload["output_0_vkformat"] == GRID_SAMPLER_2D_SAMPLER_VK_FORMAT
-    assert any(node.target == exir_ops.edge.aten.cat.default for node in nodes)
+    assert any(
+        node.target == exir_ops.edge.aten.constant_pad_nd.default for node in nodes
+    )
+    assert not any(node.target == exir_ops.edge.aten.cat.default for node in nodes)
     assert any(node.target == exir_ops.edge.aten.slice_copy.Tensor for node in nodes)
 
 
@@ -603,11 +606,17 @@ def test_rewrite_grid_sampler_to_tosa_custom_keeps_clean_graph_for_c3_sampler_pa
         == 2
     )
     assert (
-        _count_call_function_target(graph_module, exir_ops.edge.aten.cat.default) == 1
+        _count_call_function_target(
+            graph_module, exir_ops.edge.aten.constant_pad_nd.default
+        )
+        == 1
+    )
+    assert (
+        _count_call_function_target(graph_module, exir_ops.edge.aten.cat.default) == 0
     )
     assert (
         _count_call_function_target(graph_module, exir_ops.edge.aten.slice_copy.Tensor)
-        >= 2
+        == 1
     )
 
 

backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py

@@ -278,29 +278,45 @@ def _create_grid_qparam_placeholders(
     def _pad_c3_input_to_c4(
         graph_module: torch.fx.GraphModule,
         input_tensor: torch.fx.Node,
+        input_qparams: QuantArgs | None = None,
     ) -> torch.fx.Node:
         input_val = input_tensor.meta["val"]
-        first_channel = create_node(
-            graph_module.graph,
-            op_target=exir_ops.edge.aten.slice_copy.Tensor,
-            args=(input_tensor, 1, 0, 1, 1),
-            from_node=input_tensor,
-        )
-        first_channel_val = exir_ops.edge.aten.slice_copy.Tensor(input_val, 1, 0, 1, 1)
-        _set_fake_tensor_meta(first_channel, first_channel_val)
-
         padded_input = create_node(
             graph_module.graph,
-            op_target=exir_ops.edge.aten.cat.default,
-            args=([input_tensor, first_channel], 1),
+            op_target=exir_ops.edge.aten.constant_pad_nd.default,
+            args=(input_tensor, [0, 0, 0, 0, 0, 1], 0),
             from_node=input_tensor,
+            inherit_qparams=True,
         )
+        pad_qparams = input_qparams
+        if pad_qparams is None:
+            try:
+                pad_qparams = next(iter(get_output_qparams(input_tensor).values()))
+            except ValueError:
+                try:
+                    pad_qparams = next(iter(get_input_qparams(input_tensor).values()))
+                except ValueError:
+                    pad_qparams = None
+        if pad_qparams is not None:
+            padded_input.meta["input_qparams"] = {0: pad_qparams}
+            padded_input.meta["output_qparams"] = {0: pad_qparams}
         _set_fake_tensor_meta(
             padded_input,
-            exir_ops.edge.aten.cat.default([input_val, first_channel_val], 1),
+            exir_ops.edge.aten.constant_pad_nd.default(
+                input_val, [0, 0, 0, 0, 0, 1], 0
+            ),
         )
         return padded_input
 
+    @staticmethod
+    def _get_optional_input_qparams(
+        node: torch.fx.Node, input_index: int
+    ) -> QuantArgs | None:
+        try:
+            return get_input_qparams(node).get(input_index)
+        except ValueError:
+            return None
+
     @staticmethod
     def _slice_c4_output_to_c3(
         graph_module: torch.fx.GraphModule,
@@ -372,7 +388,11 @@ def call(self, graph_module):
 
             with graph_module.graph.inserting_before(node):
                 custom_input = (
-                    self._pad_c3_input_to_c4(graph_module, input_tensor)
+                    self._pad_c3_input_to_c4(
+                        graph_module,
+                        input_tensor,
+                        input_qparams=self._get_optional_input_qparams(node, 0),
+                    )
                     if pad_c3_for_sampler
                     else input_tensor
                 )

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -54,7 +54,9 @@ def __init__(self, extra_permutable_ops: set | None = None) -> None:
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.clamp.default,
             exir_ops.edge.aten.cat.default,
+            exir_ops.edge.aten.constant_pad_nd.default,
             exir_ops.edge.aten.mean.dim,
+            exir_ops.edge.aten.pad.default,
             exir_ops.edge.aten.sum.dim_IntList,
             exir_ops.edge.aten.slice_copy.Tensor,
         }
@@ -80,6 +82,11 @@ def __init__(self, extra_permutable_ops: set | None = None) -> None:
 
     _SQUEEZE_OPS = (exir_ops.edge.aten.squeeze_copy.dim,)
 
+    _PAD_OPS = (
+        exir_ops.edge.aten.constant_pad_nd.default,
+        exir_ops.edge.aten.pad.default,
+    )
+
     @staticmethod
     def _find_extra_one(longer: list[int], shorter: list[int]) -> int:
         """If longer has exactly one more element of value 1, return its index. Else -1."""
@@ -450,6 +457,8 @@ def _is_pointwise(target) -> bool:
         return False
 
     def is_node_permutable(self, node: torch.fx.Node) -> bool:
+        if node.target in self._PAD_OPS and not self._is_constant_pad(node):
+            return False
         if node.target in self._permutable_ops:
             if node.target in (
                 exir_ops.edge.aten.mean.dim,
@@ -462,6 +471,24 @@ def is_node_permutable(self, node: torch.fx.Node) -> bool:
             return True
         return self._is_pointwise(node.target)
 
+    def _is_constant_pad(self, node: torch.fx.Node) -> bool:
+        if len(node.args) < 2:
+            return False
+
+        pad = node.args[1]
+        if not isinstance(pad, (list, tuple)):
+            return False
+
+        if len(pad) % 2 != 0:
+            return False
+
+        if node.target == exir_ops.edge.aten.pad.default:
+            mode = node.args[2] if len(node.args) > 2 else node.kwargs.get("mode")
+            if mode not in (None, "constant"):
+                return False
+
+        return True
+
     def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
         # Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes.
         if not self._subgraph_edges_are_current(subgraph):
@@ -495,6 +522,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
                 self.update_mean_dim(node, node_start_perm)
             elif node.target == exir_ops.edge.aten.slice_copy.Tensor:
                 self.update_slice_copy(node, node_start_perm)
+            elif node.target in self._PAD_OPS:
+                self.update_pad(node, node_start_perm)
             elif node.target in self._VIEW_OPS:
                 self.update_view_copy(node, node_start_perm)
             elif node.target in self._UNSQUEEZE_OPS:
@@ -596,6 +625,24 @@ def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> No
         dim = get_arg(node, "dim", int)
         set_arg(node, "dim", start_permute[dim])
 
+    def update_pad(self, node: torch.fx.Node, start_permute: list[int]) -> None:
+        pad = list(cast(list[int], node.args[1]))
+        rank = len(start_permute)
+        pad_pairs = [[0, 0] for _ in range(rank)]
+        for pair_idx, pair_start in enumerate(range(0, len(pad), 2)):
+            dim = rank - 1 - pair_idx
+            pad_pairs[dim] = [pad[pair_start], pad[pair_start + 1]]
+
+        remapped_pairs = [pad_pairs[start_permute.index(dim)] for dim in range(rank)]
+        remapped_pad = []
+        for pair in reversed(remapped_pairs):
+            remapped_pad.extend(pair)
+
+        while len(remapped_pad) > 2 and remapped_pad[-2:] == [0, 0]:
+            remapped_pad = remapped_pad[:-2]
+
+        node.update_arg(1, remapped_pad)
+
     def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None:
         """Adjust view_copy shape arg after permute removal.
 

backends/transforms/test/test_permute_optimization_passes.py

@@ -810,6 +810,89 @@ def test_permute_view_add_sub_mul_view_permute(self) -> None:
             "permute_view_add_sub_mul_view_permute",
         )
 
+    def test_permute_constant_pad_nd_permute(self) -> None:
+        builder = GraphBuilder()
+        x_data = torch.randn(1, 64, 64, 3)
+        x = builder.placeholder("x", x_data)
+        p1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2])
+        )
+        pad = builder.call_operator(
+            op=exir_ops.edge.aten.constant_pad_nd.default,
+            args=(p1, [0, 0, 0, 0, 0, 1], 0.0),
+        )
+        p2 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(pad, [0, 2, 3, 1])
+        )
+        builder.output([p2])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        p = RemovePermutesAroundElementwiseOps()
+        result = cast(PassResult, p(original))
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
+        )
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.constant_pad_nd.default),
+            1,
+        )
+
+        pad_nodes = [
+            node
+            for node in result.graph_module.graph.nodes
+            if node.target == exir_ops.edge.aten.constant_pad_nd.default
+        ]
+        self.assertEqual(pad_nodes[0].args[1], [0, 1])
+        validate_numerics(
+            gm_before,
+            result.graph_module,
+            [x_data],
+            "permute_constant_pad_nd_permute",
+        )
+
+    def test_permute_aten_pad_permute(self) -> None:
+        builder = GraphBuilder()
+        x_data = torch.randn(1, 64, 64, 3)
+        x = builder.placeholder("x", x_data)
+        p1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2])
+        )
+        pad = builder.call_operator(
+            op=exir_ops.edge.aten.pad.default,
+            args=(p1, [0, 0, 0, 0, 0, 1], "constant", 0.0),
+        )
+        p2 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(pad, [0, 2, 3, 1])
+        )
+        builder.output([p2])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        p = RemovePermutesAroundElementwiseOps()
+        result = cast(PassResult, p(original))
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
+        )
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.pad.default), 1
+        )
+
+        pad_nodes = [
+            node
+            for node in result.graph_module.graph.nodes
+            if node.target == exir_ops.edge.aten.pad.default
+        ]
+        self.assertEqual(pad_nodes[0].args[1], [0, 1])
+        validate_numerics(
+            gm_before,
+            result.graph_module,
+            [x_data],
+            "permute_aten_pad_permute",
+        )
+
     def test_permute_squeeze_clamp_add_permute(self) -> None:
         """4D permute → squeeze(view) → hardtanh → add(with self) → 3D permute.
         Tests clamp + add interacting across a squeeze boundary."""