Remove permutes around fused_quant elementwise ops (#21481)

backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py

@@ -14,6 +14,9 @@
     TosaLoweringContext,
     TosaSpecification,
 )
+from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
+    RemovePermutesAroundElementwiseOps,
+)
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -47,6 +50,49 @@ def _count_nodes(graph_module: torch.fx.GraphModule, target) -> int:
     )
 
 
+def test_extra_permutable_ops_makes_op_permutable() -> None:
+    """Ops in extra_permutable_ops are permutable in the base pass."""
+
+    def build() -> torch.fx.GraphModule:
+        graph = torch.fx.Graph()
+        x = graph.placeholder("x")
+        x.meta["val"] = torch.randn(1, 3, 4, 5)
+        permute_in = graph.create_node(
+            "call_function",
+            PERMUTE_TARGET,
+            args=(x, [0, 2, 3, 1]),
+        )
+        permute_in.meta["val"] = torch.randn(1, 4, 5, 3)
+        rescale = graph.create_node(
+            "call_function",
+            RESCALE_TARGET,
+            args=(permute_in, torch.int8, [1.0], 0, 0),
+        )
+        rescale.meta["val"] = torch.randn(1, 4, 5, 3)
+        permute_out = graph.create_node(
+            "call_function",
+            PERMUTE_TARGET,
+            args=(rescale, [0, 3, 1, 2]),
+        )
+        permute_out.meta["val"] = torch.randn(1, 3, 4, 5)
+        graph.output(permute_out)
+        return torch.fx.GraphModule({}, graph)
+
+    # RESCALE is not permutable by default, so the boundary permutes stay.
+    baseline = RemovePermutesAroundElementwiseOps().call(build())
+    assert not baseline.modified
+    assert _count_nodes(baseline.graph_module, PERMUTE_TARGET) == 2
+
+    # Supplying RESCALE via extra_permutable_ops lets the region collapse.
+    with TosaLoweringContext(TOSA_INT_SPEC):
+        result = RemovePermutesAroundElementwiseOps(
+            extra_permutable_ops={RESCALE_TARGET}
+        ).call(build())
+    assert result.modified
+    assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0
+    assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1
+
+
 def test_remove_permutes_around_rescale_tosa_INT() -> None:
     graph = torch.fx.Graph()
     x = graph.placeholder("x")
@@ -140,7 +186,10 @@ def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> N
         )
 
     assert result.modified
-    assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 3
+    # The scalar constants are numel-1 (1,1,1,1): layout-invariant, so they are
+    # left wired directly with no compensating permute, and every permute in the
+    # region cancels.
+    assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0
     assert _count_nodes(result.graph_module, ERF_TARGET) == 1
 
 

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -116,6 +116,26 @@ def _check_squeeze_unsqueeze_view(self, node: torch.fx.Node) -> bool:
             return self._find_extra_one(in_shape, out_shape) != -1
         return False
 
+    def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool:
+        """True if ``node`` is a reshape whose input has at most one non-unit dim.
+
+        Flattening such a tensor -- e.g. the ``[1, C, 1, 1] -> [1, C]`` after a
+        global pool -- is permutation-invariant: every layout of the input
+        produces the identical output (the single non-unit run of elements is
+        contiguous regardless of which axis holds it). A permutation propagating
+        into it therefore simply dies, so the region can terminate here with no
+        compensating permute.
+        """
+        if node.target not in self._VIEW_OPS:
+            return False
+        inp = node.args[0]
+        assert isinstance(inp, torch.fx.Node)
+        shape = inp.meta["val"].shape
+        # Count a dim as non-unit unless it is a concrete size-1 (symbolic dims
+        # are treated as non-unit, i.e. conservatively not a sink).
+        non_unit = [d for d in shape if not (isinstance(d, int) and d == 1)]
+        return len(non_unit) <= 1
+
     def _adapt_permute_across_view(
         self, permute: list[int], node: torch.fx.Node
     ) -> list[int] | None:
@@ -321,6 +341,13 @@ def visit(  # noqa: C901
                     return False
             elif user.op == "output":
                 return False
+            elif not self._is_squeeze_unsqueeze_view(
+                user
+            ) and self._is_permutation_sink_view(user):
+                # The permutation dies at this reshape (see
+                # _is_permutation_sink_view), so terminate the region here with
+                # no compensating permute and no further downstream traversal.
+                continue
             elif not self.visit(
                 user, subgraph, processed_nodes, downstream_end, downstream_start
             ):
@@ -332,6 +359,13 @@ def visit(  # noqa: C901
                 if self.get_permutation(inp) != current_start_permute:
                     return False
                 subgraph.edges_in.add((inp, node))
+            elif inp.meta["val"].numel() == 1:
+                # A numel-1 input (per-tensor quant scale / zero_point, scalar
+                # constant, ...) is layout-invariant: it broadcasts identically
+                # under any permutation, so it needs no compensating permute and
+                # stays wired directly. Notably this keeps lifted per-tensor
+                # qparam placeholders as placeholders, which lowering requires.
+                continue
             elif self._is_constant(inp):
                 const_rank = self._get_node_rank(inp)
                 permute_rank = len(current_end_permute)

backends/transforms/targets.bzl

@@ -409,6 +409,7 @@ def define_common_targets():
         srcs = ["remove_permutes_around_elementwise_ops.py"],
         visibility = [
             "//executorch/backends/...",
+            "@EXECUTORCH_CLIENTS",
         ],
         deps = [
             ":permute_pass_utils",