From cc8c5223099fd0e3d1db2522d1f28ce5a7a5773f Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Thu, 30 Jul 2026 21:39:04 -0700 Subject: [PATCH] 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. Differential Revision: D113424191 --- ...ve_permutes_around_elementwise_tosa_ops.py | 51 ++++++++++++++++++- .../remove_permutes_around_elementwise_ops.py | 34 +++++++++++++ backends/transforms/targets.bzl | 1 + 3 files changed, 85 insertions(+), 1 deletion(-) 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 111869e7d9e..35e02b52573 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 @@ -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 diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 373f84230cb..16a6f2eb6bf 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/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) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 253a7b57bc9..6b39234ca52 100644 --- a/backends/transforms/targets.bzl +++ b/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",