Remove permutes around fused_quant elementwise ops (#21481)

Summary:

`ConvToChannelsLast` wraps every conv in `permute(NCHW->NHWC) -> conv ->
permute(NHWC->NCHW)`. When convs are joined by elementwise fused_quant ops
(residual add/mul, activations), permutes end up threaded through the
surrounding region and were never removed: the existing Cadence
`RemovePermutesAroundElementwiseOps` only recognizes aten/cadence elementwise
ops, not the SAS fused_quant ops -- and it would also wrongly treat their
lifted scale/zero_point operands as constants to be permuted.

This adds a fused_quant-aware permute-removal pass and wires it into the edge
optimization group:

- Extend the shared ExecuTorch `RemovePermutesAroundElementwiseOps` with a
  small overridable seam (`_permute_relevant_inputs`) so a subclass can hide
  operands from layout propagation. Behavior-preserving for existing users.
- New `RemovePermutesAroundFusedQuantElementwiseOps` (SAS) subclasses it,
  adding `fused_quant.add`/`mul` and the activation ops as permutable and
  exposing only their tensor operands, so the lifted scale/zero_point
  placeholders are never permuted/compensated (which would break lowering).
  Ops with per-channel qparams are skipped (not permutation-invariant).
- Replace the two Cadence permute passes (which no-op pre-Lower) in the
  optimization group with this single pass.
- Teach the shared subgraph engine about "permutation-sink" flattens: a
  `view_copy` whose input has <=1 non-unit dim (e.g. the `[1, C, 1, 1] -> [1, C]`
  after a global pool) is layout-invariant, so a permutation flowing into it
  simply dies. The region can terminate cleanly there with no compensating
  permute -- which lets the residual-block permutes collapse across the
  avgpool -> flatten -> classifier head instead of being stranded by it.

Note: fused_quant is currently SAS-specific, NOT yet a generic cross-backend
dialect, so the fused_quant knowledge deliberately stays in the SAS subclass
rather than the shared ExecuTorch pass. When fused_quant graduates to a shared
dialect, this can fold into the base pass via `extra_permutable_ops` + the
seam. (The permutation-sink flatten handling is generic and correctly lives in
the shared pass.)

On resnet18 the optimized graph goes from 83 permutes down to a single one (the
model-input boundary); every permute around the residual add/relu blocks and
across the global-pool flatten is removed.

Reviewed By: DrJessop

Differential Revision: D113424191

Author: mcremon-meta

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_val := inp.meta.get("val")) is not None and inp_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",

backends/transforms/test/test_permute_optimization_passes.py

@@ -1097,3 +1097,25 @@ def test_upstream_view_rank_mismatch_no_crash(self) -> None:
             [x_data],
             "upstream_view_rank_mismatch_no_crash",
         )
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Tests for RemovePermutesAroundElementwiseOps
+# ──────────────────────────────────────────────────────────────────────
+
+
+class RemovePermutesAroundElementwiseOpsTest(unittest.TestCase):
+    def test_no_permutes_is_noop(self) -> None:
+        """With no surrounding permutes, the pass makes no change."""
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(1, 4, 8, 8))
+        mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(x, x))
+        builder.output([mul])
+        original = builder.get_graph_module()
+
+        p = RemovePermutesAroundElementwiseOps()
+        result = cast(PassResult, p(original))
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
+        )