Skip to content

Commit 1428baa

Browse files
pssrawatfacebook-github-bot
authored andcommitted
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). Differential Revision: D130138174
1 parent 59fc134 commit 1428baa

3 files changed

Lines changed: 292 additions & 1 deletion

File tree

‎backends/transforms/remove_permutes_around_elementwise_ops.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ def __init__(
9696
exir_ops.edge.aten.pad.default,
9797
exir_ops.edge.aten.sum.dim_IntList,
9898
exir_ops.edge.aten.slice_copy.Tensor,
99+
exir_ops.edge.aten.expand_copy.default,
99100
}
100101
try:
101102
self._permutable_ops.add(
@@ -613,7 +614,11 @@ def _visit( # noqa: C901
613614
if self.get_permutation(inp) != current_start_permute:
614615
return False
615616
subgraph.edges_in.add((inp, node))
616-
elif (inp_val := inp.meta.get("val")) is not None and inp_val.numel() == 1:
617+
elif not isinstance(inp_val := inp.meta.get("val"), torch.Tensor | None):
618+
# A symbolic scalar such as sym_size may read a dim of a region
619+
# tensor, which removing the permutes would change.
620+
return False
621+
elif inp_val is not None and inp_val.numel() == 1:
617622
# A numel-1 input (per-tensor quant scale / zero_point, scalar
618623
# constant, ...) is layout-invariant: it broadcasts identically
619624
# under any permutation, so it needs no compensating permute and
@@ -844,6 +849,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901
844849
self.update_reduction_dim(node, node_start_perm)
845850
elif node.target == exir_ops.edge.aten.slice_copy.Tensor:
846851
self.update_slice_copy(node, node_start_perm)
852+
elif node.target == exir_ops.edge.aten.expand_copy.default:
853+
self.update_expand_copy(node, node_start_perm)
847854
elif node.target in self._PAD_OPS:
848855
self.update_pad(node, node_start_perm)
849856
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
10261033
dim = get_arg(node, "dim", int)
10271034
set_arg(node, "dim", start_permute[dim])
10281035

1036+
def update_expand_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None:
1037+
size = cast(list[int], get_arg(node, "size"))
1038+
# Reordering can move a -1 onto a dim a rank-growing expand adds, where
1039+
# -1 is invalid, so spell out every concrete size.
1040+
out_shape = node.meta["val"].shape
1041+
size = [
1042+
out if d == -1 and isinstance(out, int) else d
1043+
for d, out in zip(size, out_shape)
1044+
]
1045+
set_arg(node, "size", [size[start_permute.index(i)] for i in range(len(size))])
1046+
10291047
def update_pad(self, node: torch.fx.Node, start_permute: list[int]) -> None:
10301048
pad = list(cast(list[int], node.args[1]))
10311049
rank = len(start_permute)

‎backends/transforms/targets.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -851,6 +851,7 @@ def define_common_targets():
851851
"test/test_permute_optimization_passes.py",
852852
],
853853
deps = [
854+
"fbsource//third-party/pypi/parameterized:parameterized",
854855
"//caffe2:torch",
855856
"//executorch/backends/test:graph_builder",
856857
"//executorch/exir:lib",

‎backends/transforms/test/test_permute_optimization_passes.py‎

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
# pyre-unsafe
99

1010
import copy
11+
import itertools
1112
import unittest
1213
from typing import cast
1314

@@ -38,6 +39,7 @@
3839
from executorch.exir import EdgeCompileConfig, to_edge
3940
from executorch.exir.dialects._ops import ops as exir_ops
4041
from executorch.exir.pass_base import PassResult
42+
from parameterized import parameterized
4143
from torch.utils import _pytree as pytree
4244

4345

@@ -2734,6 +2736,276 @@ def test_lower_rank_constant_is_transposed_into_region_layout(self) -> None:
27342736
)
27352737

27362738

2739+
_NHWC_SHAPE = [2, 3, 4, 5]
2740+
_AXIS_NAMES = "NHWC"
2741+
# x.permute(_NHWC_FROM_NCHW) turns an NCHW tensor into NHWC.
2742+
_NHWC_FROM_NCHW = [0, 2, 3, 1]
2743+
2744+
2745+
def _nchw_order(nhwc: list[int]) -> list[int]:
2746+
return [nhwc[0], nhwc[3], nhwc[1], nhwc[2]]
2747+
2748+
2749+
def _case_name(broadcast: tuple[int, ...], kept_states: dict[int, str]) -> str:
2750+
parts = (
2751+
["broadcast_" + "".join(_AXIS_NAMES[a] for a in broadcast)] if broadcast else []
2752+
)
2753+
parts += [f"{_AXIS_NAMES[a]}_{state}" for a, state in kept_states.items()]
2754+
return "_".join(parts)
2755+
2756+
2757+
def _region_expand_cases() -> list[tuple[str, tuple[int, ...], dict[int, str]]]:
2758+
cases = []
2759+
for count in range(1, 5):
2760+
for broadcast in itertools.combinations(range(4), count):
2761+
kept = [a for a in range(4) if a not in broadcast]
2762+
for states in itertools.product(
2763+
("size", "minus_one", "dynamic"), repeat=len(kept)
2764+
):
2765+
kept_states = dict(zip(kept, states))
2766+
cases.append(
2767+
(_case_name(broadcast, kept_states), broadcast, kept_states)
2768+
)
2769+
return cases
2770+
2771+
2772+
def _constant_expand_cases() -> list[tuple[str, int, tuple[int, ...], dict[int, str]]]:
2773+
cases = []
2774+
for rank in range(1, 4):
2775+
own_axes = list(range(4 - rank, 4))
2776+
for count in range(rank + 1):
2777+
for broadcast in itertools.combinations(own_axes, count):
2778+
kept = [a for a in own_axes if a not in broadcast]
2779+
for states in itertools.product(
2780+
("size", "minus_one"), repeat=len(kept)
2781+
):
2782+
kept_states = dict(zip(kept, states))
2783+
name = f"rank{rank}_" + _case_name(broadcast, kept_states)
2784+
cases.append((name.rstrip("_"), rank, broadcast, kept_states))
2785+
return cases
2786+
2787+
2788+
def _export_to_edge(
2789+
module: torch.nn.Module,
2790+
example_inputs: tuple[torch.Tensor, ...],
2791+
dynamic_shapes: dict[str, dict[int, torch.export.Dim] | None] | None,
2792+
) -> torch.fx.GraphModule:
2793+
exported = torch.export.export(
2794+
module, example_inputs, dynamic_shapes=dynamic_shapes
2795+
)
2796+
edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False))
2797+
return edge.exported_program().graph_module
2798+
2799+
2800+
class RemovePermutesAroundExpandCopyTest(unittest.TestCase):
2801+
"""An NHWC region around ``expand_copy``, for every set of broadcast axes,
2802+
with kept axes written as their size, as -1, or as -1 on a dynamic dim."""
2803+
2804+
def _remove_permutes(self, original: torch.fx.GraphModule) -> torch.fx.GraphModule:
2805+
result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original))
2806+
self.assertTrue(result.modified)
2807+
gm = result.graph_module
2808+
# A constant whose non-unit dims change order keeps a permute of its
2809+
# own; only the activation path has to be permute-free.
2810+
permute = exir_ops.edge.aten.permute_copy.default
2811+
(x,) = (n for n in gm.graph.nodes if n.op == "placeholder" and n.name == "x")
2812+
self.assertNotIn(permute, [user.target for user in x.users])
2813+
(output,) = gm.graph.find_nodes(op="output")
2814+
self.assertNotEqual(output.args[0][0].target, permute)
2815+
return gm
2816+
2817+
def _expand_size(self, gm: torch.fx.GraphModule) -> list[int]:
2818+
(expand,) = gm.graph.find_nodes(
2819+
op="call_function", target=exir_ops.edge.aten.expand_copy.default
2820+
)
2821+
return expand.args[1]
2822+
2823+
@parameterized.expand(_region_expand_cases())
2824+
def test_expand_of_region_tensor(
2825+
self, name: str, broadcast: tuple[int, ...], kept_states: dict[int, str]
2826+
) -> None:
2827+
dynamic_axes = [a for a, state in kept_states.items() if state == "dynamic"]
2828+
size = [
2829+
-1 if kept_states.get(a) in ("minus_one", "dynamic") else _NHWC_SHAPE[a]
2830+
for a in range(4)
2831+
]
2832+
2833+
class PermuteExpandPermute(torch.nn.Module):
2834+
def forward(self, x: torch.Tensor) -> torch.Tensor:
2835+
return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2)
2836+
2837+
input_nhwc = [1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4)]
2838+
x = torch.randn(*_nchw_order(input_nhwc))
2839+
dims = {
2840+
_NHWC_FROM_NCHW[a]: torch.export.Dim(f"d{a}", min=2, max=16)
2841+
for a in dynamic_axes
2842+
}
2843+
original = _export_to_edge(PermuteExpandPermute(), (x,), {"x": dims or None})
2844+
gm_before = copy.deepcopy(original)
2845+
2846+
gm = self._remove_permutes(original)
2847+
2848+
expected_nhwc = [-1 if a in dynamic_axes else _NHWC_SHAPE[a] for a in range(4)]
2849+
self.assertEqual(self._expand_size(gm), _nchw_order(expected_nhwc))
2850+
resized_nhwc = [
2851+
d + 3 if a in dynamic_axes else d for a, d in enumerate(input_nhwc)
2852+
]
2853+
for inputs in ([x], [torch.randn(*_nchw_order(resized_nhwc))]):
2854+
validate_numerics(gm_before, gm, inputs, name)
2855+
2856+
@parameterized.expand(_constant_expand_cases())
2857+
def test_rank_growing_expand_of_constant(
2858+
self,
2859+
name: str,
2860+
rank: int,
2861+
broadcast: tuple[int, ...],
2862+
kept_states: dict[int, str],
2863+
) -> None:
2864+
bias_shape = [
2865+
1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4 - rank, 4)
2866+
]
2867+
size = [
2868+
-1 if kept_states.get(a) == "minus_one" else _NHWC_SHAPE[a]
2869+
for a in range(4)
2870+
]
2871+
builder = GraphBuilder()
2872+
x_data = torch.randn(*_nchw_order(_NHWC_SHAPE))
2873+
bias_data = torch.randn(*bias_shape)
2874+
x = builder.placeholder("x", x_data)
2875+
bias = builder.placeholder("p_bias", bias_data)
2876+
nhwc = builder.call_operator(
2877+
op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1])
2878+
)
2879+
expanded = builder.call_operator(
2880+
op=exir_ops.edge.aten.expand_copy.default, args=(bias, size)
2881+
)
2882+
scaled = builder.call_operator(
2883+
op=exir_ops.edge.aten.mul.Tensor, args=(nhwc, expanded)
2884+
)
2885+
nchw = builder.call_operator(
2886+
op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [0, 3, 1, 2])
2887+
)
2888+
builder.output([nchw])
2889+
original = builder.get_graph_module()
2890+
gm_before = copy.deepcopy(original)
2891+
2892+
gm = self._remove_permutes(original)
2893+
2894+
self.assertEqual(self._expand_size(gm), _nchw_order(_NHWC_SHAPE))
2895+
validate_numerics(gm_before, gm, [x_data, bias_data], name)
2896+
2897+
def test_keep_entry_moved_onto_an_added_dim(self) -> None:
2898+
# Undoing the [1, 0, 2, 3] swap moves the -1 onto the dim the expand
2899+
# adds, so it has to be spelled out as 1 first.
2900+
builder = GraphBuilder()
2901+
x_data = torch.randn(1, 4, 5, 3)
2902+
bias_data = torch.randn(1, 5, 3)
2903+
x = builder.placeholder("x", x_data)
2904+
bias = builder.placeholder("p_bias", bias_data)
2905+
swapped = builder.call_operator(
2906+
op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 0, 2, 3])
2907+
)
2908+
expanded = builder.call_operator(
2909+
op=exir_ops.edge.aten.expand_copy.default, args=(bias, [4, -1, 5, 3])
2910+
)
2911+
scaled = builder.call_operator(
2912+
op=exir_ops.edge.aten.mul.Tensor, args=(swapped, expanded)
2913+
)
2914+
back = builder.call_operator(
2915+
op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [1, 0, 2, 3])
2916+
)
2917+
builder.output([back])
2918+
original = builder.get_graph_module()
2919+
gm_before = copy.deepcopy(original)
2920+
2921+
gm = self._remove_permutes(original)
2922+
2923+
self.assertEqual(self._expand_size(gm), [1, 4, 5, 3])
2924+
validate_numerics(gm_before, gm, [x_data, bias_data], "added_dim")
2925+
2926+
@parameterized.expand([(f"symbolic_target_{_AXIS_NAMES[a]}", a) for a in range(4)])
2927+
def test_symbolic_broadcast_target_leaves_region(
2928+
self, name: str, axis: int
2929+
) -> None:
2930+
class ExpandToSymbolicSize(torch.nn.Module):
2931+
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
2932+
size = [-1] * 4
2933+
size[axis] = t.shape[0]
2934+
return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2)
2935+
2936+
input_nhwc = [1 if a == axis else _NHWC_SHAPE[a] for a in range(4)]
2937+
x = torch.randn(*_nchw_order(input_nhwc))
2938+
original = _export_to_edge(
2939+
ExpandToSymbolicSize(),
2940+
(x, torch.randn(6)),
2941+
{"x": None, "t": {0: torch.export.Dim("t", min=2, max=16)}},
2942+
)
2943+
gm_before = copy.deepcopy(original)
2944+
2945+
result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original))
2946+
2947+
self.assertFalse(result.modified)
2948+
self.assertEqual(
2949+
count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2
2950+
)
2951+
for t_size in (6, 9):
2952+
validate_numerics(
2953+
gm_before, result.graph_module, [x, torch.randn(t_size)], name
2954+
)
2955+
2956+
2957+
class SymbolicScalarInputTest(unittest.TestCase):
2958+
"""A region whose ops take a symbolic size is left alone rather than
2959+
crashing the pass."""
2960+
2961+
def test_symbolic_arguments_leave_the_region_unchanged(self) -> None:
2962+
def nhwc(x: torch.Tensor) -> torch.Tensor:
2963+
return x.permute(0, 2, 3, 1)
2964+
2965+
def nchw(x: torch.Tensor) -> torch.Tensor:
2966+
return x.permute(0, 3, 1, 2)
2967+
2968+
class SymbolicArgument(torch.nn.Module):
2969+
def __init__(self, op: str) -> None:
2970+
super().__init__()
2971+
self.op = op
2972+
2973+
def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
2974+
y = nhwc(x)
2975+
if self.op == "expand":
2976+
y = y.expand(2, n.shape[0], 8, 3)
2977+
elif self.op == "slice":
2978+
y = y[:, :, : n.shape[0], :]
2979+
else:
2980+
y = torch.nn.functional.pad(y, (0, 0, 0, n.shape[0]))
2981+
return nchw(y)
2982+
2983+
for op in ("expand", "slice", "pad"):
2984+
with self.subTest(op=op):
2985+
x = torch.randn(2, 3, 1 if op == "expand" else 4, 8)
2986+
original = _export_to_edge(
2987+
SymbolicArgument(op),
2988+
(x, torch.randn(4)),
2989+
{"x": None, "n": {0: torch.export.Dim("n", min=2, max=7)}},
2990+
)
2991+
gm_before = copy.deepcopy(original)
2992+
2993+
result = cast(
2994+
PassResult, RemovePermutesAroundElementwiseOps()(original)
2995+
)
2996+
2997+
self.assertFalse(result.modified)
2998+
self.assertEqual(
2999+
count_node(
3000+
result.graph_module, exir_ops.edge.aten.permute_copy.default
3001+
),
3002+
2,
3003+
)
3004+
validate_numerics(
3005+
gm_before, result.graph_module, [x, torch.randn(6)], op
3006+
)
3007+
3008+
27373009
class LayoutPermuteVisibilityTest(unittest.TestCase):
27383010
"""The data-movement passes must see both permute dialects.
27393011

0 commit comments

Comments
 (0)