From c78343617e30f8ae7aa079cac97e15d0e9f9b50b Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Mon, 9 Nov 2026 13:44:58 -0800 Subject: [PATCH] Remove permutes around expand_copy; reject symbolic scalar inputs (#24742) Summary: `RemovePermutesAroundElementwiseOps` stopped at `aten.expand_copy.default`, so a layout region containing a broadcast was cut in two and kept a permute pair on each side. A squeeze-and-excitation gate, a per-channel scale broadcast over the time axis and multiplied back in, puts an expand in the middle of exactly such a region. `expand_copy` is now in the default permutable set, and `update_expand_copy` un-permutes its size argument the same way the other `update_*` methods remap their dims. Before reordering, it replaces each `-1` with the concrete output size, keeping `-1` only where that dim is symbolic: in a rank-growing expand, reordering could otherwise move a `-1` onto a dim the expand adds, where `-1` is invalid. With that, rank-growing expands of constants are remapped correctly and need no filter; non-constant lower-rank inputs are still rejected by the region's rank check. The pass also crashed with `AttributeError: 'SymInt' object has no attribute 'numel'` when a region op took a symbolic argument, such as `slice_copy` with an end of `y.shape[0]`, `constant_pad_nd` with a symbolic pad amount, or `expand_copy` with a symbolic size. The upstream walk reaches the `sym_size` node and calls `.numel()` on its `SymInt` value. Such inputs now reject the region. They cannot be skipped like numel-1 tensors: a `sym_size` can read a dim of a region tensor, which removing the permutes would change. Authored with an AI coding assistant (Claude). Reviewed By: mcremon-meta Differential Revision: D130138174 --- .../remove_permutes_around_elementwise_ops.py | 20 +- backends/transforms/targets.bzl | 1 + .../test/test_permute_optimization_passes.py | 272 ++++++++++++++++++ 3 files changed, 292 insertions(+), 1 deletion(-) diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index b237b730882..c80e4723d9d 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -96,6 +96,7 @@ def __init__( exir_ops.edge.aten.pad.default, exir_ops.edge.aten.sum.dim_IntList, exir_ops.edge.aten.slice_copy.Tensor, + exir_ops.edge.aten.expand_copy.default, } try: self._permutable_ops.add( @@ -613,7 +614,11 @@ 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: + elif not isinstance(inp_val := inp.meta.get("val"), torch.Tensor | None): + # A symbolic scalar such as sym_size may read a dim of a region + # tensor, which removing the permutes would change. + return False + elif inp_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 @@ -844,6 +849,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 self.update_reduction_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 == exir_ops.edge.aten.expand_copy.default: + self.update_expand_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: @@ -1026,6 +1033,17 @@ 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_expand_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: + size = cast(list[int], get_arg(node, "size")) + # Reordering can move a -1 onto a dim a rank-growing expand adds, where + # -1 is invalid, so spell out every concrete size. + out_shape = node.meta["val"].shape + size = [ + out if d == -1 and isinstance(out, int) else d + for d, out in zip(size, out_shape) + ] + set_arg(node, "size", [size[start_permute.index(i)] for i in range(len(size))]) + 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) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 88142e33fac..05178181286 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -852,6 +852,7 @@ def define_common_targets(): "test/test_permute_optimization_passes.py", ], deps = [ + "fbsource//third-party/pypi/parameterized:parameterized", "//caffe2:torch", "//executorch/backends/test:graph_builder", "//executorch/exir:lib", diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 5f9fae482f6..8b87c5712ce 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -8,6 +8,7 @@ # pyre-unsafe import copy +import itertools import unittest from typing import cast @@ -38,6 +39,7 @@ from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import PassResult +from parameterized import parameterized from torch.utils import _pytree as pytree @@ -2734,6 +2736,276 @@ def test_lower_rank_constant_is_transposed_into_region_layout(self) -> None: ) +_NHWC_SHAPE = [2, 3, 4, 5] +_AXIS_NAMES = "NHWC" +# x.permute(_NHWC_FROM_NCHW) turns an NCHW tensor into NHWC. +_NHWC_FROM_NCHW = [0, 2, 3, 1] + + +def _nchw_order(nhwc: list[int]) -> list[int]: + return [nhwc[0], nhwc[3], nhwc[1], nhwc[2]] + + +def _case_name(broadcast: tuple[int, ...], kept_states: dict[int, str]) -> str: + parts = ( + ["broadcast_" + "".join(_AXIS_NAMES[a] for a in broadcast)] if broadcast else [] + ) + parts += [f"{_AXIS_NAMES[a]}_{state}" for a, state in kept_states.items()] + return "_".join(parts) + + +def _region_expand_cases() -> list[tuple[str, tuple[int, ...], dict[int, str]]]: + cases = [] + for count in range(1, 5): + for broadcast in itertools.combinations(range(4), count): + kept = [a for a in range(4) if a not in broadcast] + for states in itertools.product( + ("size", "minus_one", "dynamic"), repeat=len(kept) + ): + kept_states = dict(zip(kept, states)) + cases.append( + (_case_name(broadcast, kept_states), broadcast, kept_states) + ) + return cases + + +def _constant_expand_cases() -> list[tuple[str, int, tuple[int, ...], dict[int, str]]]: + cases = [] + for rank in range(1, 4): + own_axes = list(range(4 - rank, 4)) + for count in range(rank + 1): + for broadcast in itertools.combinations(own_axes, count): + kept = [a for a in own_axes if a not in broadcast] + for states in itertools.product( + ("size", "minus_one"), repeat=len(kept) + ): + kept_states = dict(zip(kept, states)) + name = f"rank{rank}_" + _case_name(broadcast, kept_states) + cases.append((name.rstrip("_"), rank, broadcast, kept_states)) + return cases + + +def _export_to_edge( + module: torch.nn.Module, + example_inputs: tuple[torch.Tensor, ...], + dynamic_shapes: dict[str, dict[int, torch.export.Dim] | None] | None, +) -> torch.fx.GraphModule: + exported = torch.export.export( + module, example_inputs, dynamic_shapes=dynamic_shapes + ) + edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False)) + return edge.exported_program().graph_module + + +class RemovePermutesAroundExpandCopyTest(unittest.TestCase): + """An NHWC region around ``expand_copy``, for every set of broadcast axes, + with kept axes written as their size, as -1, or as -1 on a dynamic dim.""" + + def _remove_permutes(self, original: torch.fx.GraphModule) -> torch.fx.GraphModule: + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertTrue(result.modified) + gm = result.graph_module + # A constant whose non-unit dims change order keeps a permute of its + # own; only the activation path has to be permute-free. + permute = exir_ops.edge.aten.permute_copy.default + (x,) = (n for n in gm.graph.nodes if n.op == "placeholder" and n.name == "x") + self.assertNotIn(permute, [user.target for user in x.users]) + (output,) = gm.graph.find_nodes(op="output") + self.assertNotEqual(output.args[0][0].target, permute) + return gm + + def _expand_size(self, gm: torch.fx.GraphModule) -> list[int]: + (expand,) = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.expand_copy.default + ) + return expand.args[1] + + @parameterized.expand(_region_expand_cases()) + def test_expand_of_region_tensor( + self, name: str, broadcast: tuple[int, ...], kept_states: dict[int, str] + ) -> None: + dynamic_axes = [a for a, state in kept_states.items() if state == "dynamic"] + size = [ + -1 if kept_states.get(a) in ("minus_one", "dynamic") else _NHWC_SHAPE[a] + for a in range(4) + ] + + class PermuteExpandPermute(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2) + + input_nhwc = [1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4)] + x = torch.randn(*_nchw_order(input_nhwc)) + dims = { + _NHWC_FROM_NCHW[a]: torch.export.Dim(f"d{a}", min=2, max=16) + for a in dynamic_axes + } + original = _export_to_edge(PermuteExpandPermute(), (x,), {"x": dims or None}) + gm_before = copy.deepcopy(original) + + gm = self._remove_permutes(original) + + expected_nhwc = [-1 if a in dynamic_axes else _NHWC_SHAPE[a] for a in range(4)] + self.assertEqual(self._expand_size(gm), _nchw_order(expected_nhwc)) + resized_nhwc = [ + d + 3 if a in dynamic_axes else d for a, d in enumerate(input_nhwc) + ] + for inputs in ([x], [torch.randn(*_nchw_order(resized_nhwc))]): + validate_numerics(gm_before, gm, inputs, name) + + @parameterized.expand(_constant_expand_cases()) + def test_rank_growing_expand_of_constant( + self, + name: str, + rank: int, + broadcast: tuple[int, ...], + kept_states: dict[int, str], + ) -> None: + bias_shape = [ + 1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4 - rank, 4) + ] + size = [ + -1 if kept_states.get(a) == "minus_one" else _NHWC_SHAPE[a] + for a in range(4) + ] + builder = GraphBuilder() + x_data = torch.randn(*_nchw_order(_NHWC_SHAPE)) + bias_data = torch.randn(*bias_shape) + x = builder.placeholder("x", x_data) + bias = builder.placeholder("p_bias", bias_data) + nhwc = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) + ) + expanded = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, args=(bias, size) + ) + scaled = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(nhwc, expanded) + ) + nchw = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [0, 3, 1, 2]) + ) + builder.output([nchw]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + gm = self._remove_permutes(original) + + self.assertEqual(self._expand_size(gm), _nchw_order(_NHWC_SHAPE)) + validate_numerics(gm_before, gm, [x_data, bias_data], name) + + def test_keep_entry_moved_onto_an_added_dim(self) -> None: + # Undoing the [1, 0, 2, 3] swap moves the -1 onto the dim the expand + # adds, so it has to be spelled out as 1 first. + builder = GraphBuilder() + x_data = torch.randn(1, 4, 5, 3) + bias_data = torch.randn(1, 5, 3) + x = builder.placeholder("x", x_data) + bias = builder.placeholder("p_bias", bias_data) + swapped = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 0, 2, 3]) + ) + expanded = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, args=(bias, [4, -1, 5, 3]) + ) + scaled = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(swapped, expanded) + ) + back = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [1, 0, 2, 3]) + ) + builder.output([back]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + gm = self._remove_permutes(original) + + self.assertEqual(self._expand_size(gm), [1, 4, 5, 3]) + validate_numerics(gm_before, gm, [x_data, bias_data], "added_dim") + + @parameterized.expand([(f"symbolic_target_{_AXIS_NAMES[a]}", a) for a in range(4)]) + def test_symbolic_broadcast_target_leaves_region( + self, name: str, axis: int + ) -> None: + class ExpandToSymbolicSize(torch.nn.Module): + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + size = [-1] * 4 + size[axis] = t.shape[0] + return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2) + + input_nhwc = [1 if a == axis else _NHWC_SHAPE[a] for a in range(4)] + x = torch.randn(*_nchw_order(input_nhwc)) + original = _export_to_edge( + ExpandToSymbolicSize(), + (x, torch.randn(6)), + {"x": None, "t": {0: torch.export.Dim("t", min=2, max=16)}}, + ) + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + for t_size in (6, 9): + validate_numerics( + gm_before, result.graph_module, [x, torch.randn(t_size)], name + ) + + +class SymbolicScalarInputTest(unittest.TestCase): + """A region whose ops take a symbolic size is left alone rather than + crashing the pass.""" + + def test_symbolic_arguments_leave_the_region_unchanged(self) -> None: + def nhwc(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 3, 1) + + def nchw(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 3, 1, 2) + + class SymbolicArgument(torch.nn.Module): + def __init__(self, op: str) -> None: + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor: + y = nhwc(x) + if self.op == "expand": + y = y.expand(2, n.shape[0], 8, 3) + elif self.op == "slice": + y = y[:, :, : n.shape[0], :] + else: + y = torch.nn.functional.pad(y, (0, 0, 0, n.shape[0])) + return nchw(y) + + for op in ("expand", "slice", "pad"): + with self.subTest(op=op): + x = torch.randn(2, 3, 1 if op == "expand" else 4, 8) + original = _export_to_edge( + SymbolicArgument(op), + (x, torch.randn(4)), + {"x": None, "n": {0: torch.export.Dim("n", min=2, max=7)}}, + ) + gm_before = copy.deepcopy(original) + + result = cast( + PassResult, RemovePermutesAroundElementwiseOps()(original) + ) + + self.assertFalse(result.modified) + self.assertEqual( + count_node( + result.graph_module, exir_ops.edge.aten.permute_copy.default + ), + 2, + ) + validate_numerics( + gm_before, result.graph_module, [x, torch.randn(6)], op + ) + + class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects.