From 2ec948f6edb97523e5e2f870369c8481cd7493ab Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Wed, 12 Aug 2026 16:52:32 -0700 Subject: [PATCH] Recognise the repeat_interleave idiom in RemovePermutesAroundElementwiseOps (#21790) Summary: `repeat_interleave(scale, dim)` lowers to a three-node idiom:: unsqueeze(dim + 1) -> expand_copy(scale at dim + 1) -> view_copy(merge dim, dim + 1) which the region walker could not cross, so any permuted region containing one was rejected. torchaudio's `Stretch2d` is built from this, so it blocks every upsampling network -- on wavernn it strands the permutes `ConvToChannelsLast` put around each conv in the upsample path. Adding `expand_copy` to `_permutable_ops` on its own does not help: traversal then dies one node later at the merging `view_copy`. That merge is *not* unconditionally layout-invariant the way `cat`/`slice` are -- flattening two dims only commutes with a permutation that leaves them adjacent and in order, otherwise the flatten reorders elements. Worse, handling the three nodes separately would force a choice of un-permuted position for the intermediate unit dim (inserting a size-1 dim is ambiguous: any position gives a consistent permutation). That arbitrary choice then decides whether the later merge is still legal, coupling two rewrites that look independent. So the triple is matched as a unit (`_interleave_triple`). It is rank-preserving overall, so the permutation flows through unchanged and the only rewrite needed is remapping the interleaved dim, `new_dim = P[old_dim]`, plus rebuilding the three shape args in un-permuted space (`update_interleave`). No intermediate layout is ever materialised, so the ambiguity does not arise. Guards: head and expand must each have exactly one user (otherwise the in-place rewrite would corrupt another consumer), every non-inserted dim of the expand must pass through untouched, and the view's output shape must be exactly the merge. Symbolic shapes are declined. The head may be either an explicit `unsqueeze_copy` or a `view_copy` that inserts one unit dim, matching how the pass already treats the two spellings interchangeably. Differential Revision: D114508262 --- .../remove_permutes_around_elementwise_ops.py | 267 +++++++++++++- .../test/test_permute_optimization_passes.py | 329 ++++++++++++++++++ 2 files changed, 577 insertions(+), 19 deletions(-) diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 65b0c6c2e59..6e916dfe50a 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -23,6 +23,8 @@ class RemovePermutesAroundElementwiseOps(ExportPass): permutes if possible. Allows special handling for certain non-elementwise ops that can be easily updated based on the permute's parameter such as mean, cat, and slice. + The repeat_interleave idiom (unsqueeze -> expand_copy -> merging view_copy) is + recognised as a single rank-preserving unit; see _interleave_triple. """ @dataclass() @@ -44,6 +46,11 @@ class Subgraph: node_end_permute: dict[torch.fx.Node, list[int]] = field(default_factory=dict) # Per-node expected start permutation for upstream traversal. node_start_permute: dict[torch.fx.Node, list[int]] = field(default_factory=dict) + # repeat_interleave triples keyed by their unit-dim-inserting head node, + # mapping to (dim, scale, expand_node, view_node). See _interleave_triple. + interleaves: dict[ + torch.fx.Node, tuple[int, int, torch.fx.Node, torch.fx.Node] + ] = field(default_factory=dict) def __init__(self, extra_permutable_ops: set | None = None) -> None: super().__init__() @@ -72,6 +79,10 @@ def __init__(self, extra_permutable_ops: set | None = None) -> None: if extra_permutable_ops: self._permutable_ops |= extra_permutable_ops self._sq_unsq_cache: dict[torch.fx.Node, bool] = {} + self._interleave_cache: dict[ + torch.fx.Node, + tuple[int, int, torch.fx.Node, torch.fx.Node] | None, + ] = {} _VIEW_OPS = ( exir_ops.edge.aten.view_copy.default, @@ -165,6 +176,101 @@ def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool: non_unit = [d for d in shape if not (isinstance(d, int) and d == 1)] return len(non_unit) <= 1 + def _inserted_unit_dim(self, node: torch.fx.Node) -> int | None: + """Position of the size-1 dim ``node`` inserts, else None. + + Accepts both an explicit unsqueeze and a view_copy that only adds a + single unit dim, matching how the rest of the pass treats the two + spellings interchangeably. + """ + is_unsqueeze = node.target == exir_ops.edge.aten.unsqueeze_copy.default + if not is_unsqueeze and node.target not in self._VIEW_OPS: + return None + inp = node.args[0] + if not isinstance(inp, torch.fx.Node): + return None + in_shape = self._concrete_shape(inp) + out_shape = self._concrete_shape(node) + if in_shape is None or out_shape is None: + return None + if len(out_shape) != len(in_shape) + 1: + return None + if is_unsqueeze: + dim = get_arg(node, "dim", int) + pos = dim if dim >= 0 else dim + len(out_shape) + if not 0 <= pos < len(out_shape) or out_shape[pos] != 1: + return None + return pos + positions = self._find_extra_ones(out_shape, in_shape) + if positions is None or len(positions) != 1: + return None + return positions[0] + + def _interleave_triple( + self, node: torch.fx.Node + ) -> tuple[int, int, torch.fx.Node, torch.fx.Node] | None: + """Recognise a repeat_interleave and return (dim, scale, expand, view).""" + if node not in self._interleave_cache: + self._interleave_cache[node] = self._match_interleave_triple(node) + return self._interleave_cache[node] + + def _match_interleave_triple( + self, node: torch.fx.Node + ) -> tuple[int, int, torch.fx.Node, torch.fx.Node] | None: + """Match a repeat_interleave lowered to three shape operations. + + ``repeat_interleave(scale, dim)`` lowers to:: + + unsqueeze(dim + 1) -> expand_copy(scale at dim + 1) + -> view_copy(merge dim, dim + 1) + + (e.g. torchaudio's Stretch2d). The triple is rank-preserving overall, so + a permutation flows through it unchanged and only the dim it acts on has + to be remapped -- unlike the merging view_copy on its own, which is not + layout-invariant. Handling the three nodes as one unit also avoids having + to pick an un-permuted position for the intermediate unit dim, a choice + that would otherwise decide whether the merge stays legal. + """ + pos = self._inserted_unit_dim(node) + if pos is None or pos == 0 or len(node.users) != 1: + return None + dim = pos - 1 + + expand_node = next(iter(node.users)) + if ( + expand_node.target != exir_ops.edge.aten.expand_copy.default + or len(expand_node.users) != 1 + ): + return None + + unsq_shape = self._concrete_shape(node) + if unsq_shape is None: + return None + size = get_arg(expand_node, "size") + if not isinstance(size, (list, tuple)) or len(size) != len(unsq_shape): + return None + size = list(size) + if not all(isinstance(s, int) for s in size): + return None + # Every dim other than the inserted one must pass through untouched. + if any(s != -1 and s != unsq_shape[k] for k, s in enumerate(size) if k != pos): + return None + scale = size[pos] + if scale < 1: + return None + + view_node = next(iter(expand_node.users)) + if view_node.target not in self._VIEW_OPS: + return None + in_shape = self._concrete_shape(cast(torch.fx.Node, node.args[0])) + if in_shape is None: + return None + merged = list(in_shape) + merged[dim] *= scale + if self._concrete_shape(view_node) != merged: + return None + return dim, scale, expand_node, view_node + def _adapt_permute_across_view( self, permute: list[int], node: torch.fx.Node ) -> list[int] | None: @@ -216,6 +322,7 @@ def _adapt_permute_across_view( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 self._sq_unsq_cache.clear() + self._interleave_cache.clear() subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = [] processed_nodes: set[torch.fx.Node] = set() for node in graph_module.graph.find_nodes( @@ -229,7 +336,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Try direct users first (same-rank matching) for user in node.users: - if not self.is_node_permutable(user): + if ( + not self.is_node_permutable(user) + and self._interleave_triple(user) is None + ): continue subgraph = self.Subgraph(start_permute, end_permute) if self.visit(user, subgraph, processed_nodes): @@ -256,7 +366,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 adapted_start.index(i) for i in range(len(adapted_start)) ] for view_user in view_node.users: - if not self.is_node_permutable(view_user): + if ( + not self.is_node_permutable(view_user) + and self._interleave_triple(view_user) is None + ): continue subgraph = self.Subgraph(adapted_start, adapted_end) # Include the view in the subgraph @@ -305,29 +418,56 @@ def visit( # noqa: C901 if node in subgraph.nodes: return True - if node in processed_nodes or not self.is_node_permutable(node): + if node in processed_nodes: + return False + # Explicit unsqueeze nodes are not generally permutable after shape-op + # canonicalization, but a guarded repeat_interleave triple is handled as + # one rank-preserving unit. + triple = self._interleave_triple(node) + if triple is None and not self.is_node_permutable(node): return False - # A permutable op can still change rank via broadcasting (e.g. - # [8, 1] + [1, 1, 1] -> [1, 8, 1]), which would leave the node carrying - # a permutation of the wrong rank. Downstream rewrites index the - # permutation by dim (update_cat / update_mean_dim / update_slice_copy), - # so bail out rather than mis-permute or index out of range. - # Squeeze/unsqueeze views are exempt: they intentionally carry their - # input-rank permutation and are rank-checked in - # _adapt_permute_across_view. - if not self._is_squeeze_unsqueeze_view(node): - node_shape = getattr(node.meta.get("val"), "shape", None) - if node_shape is not None and len(node_shape) != len(current_start_permute): + if triple is not None: + inp = node.args[0] if node.args else None + if not isinstance(inp, torch.fx.Node): return False + in_shape = self._concrete_shape(inp) + if in_shape is None or len(current_start_permute) != len(in_shape): + return False + else: + # A permutable op can still change rank via broadcasting (e.g. + # [8, 1] + [1, 1, 1] -> [1, 8, 1]), which would leave the node + # carrying a permutation of the wrong rank. Squeeze/unsqueeze views + # are exempt because _adapt_permute_across_view checks their ranks. + if not self._is_squeeze_unsqueeze_view(node): + node_shape = getattr(node.meta.get("val"), "shape", None) + if node_shape is not None and len(node_shape) != len( + current_start_permute + ): + return False subgraph.nodes.add(node) subgraph.node_end_permute[node] = current_end_permute subgraph.node_start_permute[node] = current_start_permute + # A repeat_interleave triple is absorbed whole: its interior nodes are + # not layout-invariant individually, but the triple is rank-preserving. + users_source = node + if triple is not None: + users_source = self._absorb_interleave( + node, + triple, + subgraph, + processed_nodes, + current_end_permute, + current_start_permute, + ) + if users_source is None: + return False + # If this is a squeeze/unsqueeze view, adapt permutations for # traversal across the rank change boundary. downstream_end = current_end_permute downstream_start = current_start_permute - if self._is_squeeze_unsqueeze_view(node): + if triple is None and self._is_squeeze_unsqueeze_view(node): # Adapt start permute for downstream (input-rank → output-rank) adapted_start = self._adapt_permute_across_view(current_start_permute, node) if adapted_start is None: @@ -342,11 +482,11 @@ def visit( # noqa: C901 downstream_end = [adapted_start.index(i) for i in range(len(adapted_start))] # Traverse downstream: - for user in node.users: + for user in users_source.users: if user.target == exir_ops.edge.aten.permute_copy.default: user_perm = self.get_permutation(user) if user_perm == downstream_end: - subgraph.edges_out.add((node, user)) + subgraph.edges_out.add((users_source, user)) else: # Check if permute → view(squeeze/unsqueeze) forms an # end boundary at a different rank. @@ -354,7 +494,7 @@ def visit( # noqa: C901 if len(user_users) == 1 and self._is_squeeze_unsqueeze_view( user_users[0] ): - view_after = user_users[0] + view_after: torch.fx.Node = user_users[0] # Adapt the start permute across the view and derive # the expected end permute as its inverse. adapted_start_after = self._adapt_permute_across_view( @@ -367,7 +507,7 @@ def visit( # noqa: C901 ] if user_perm == adapted: # Include both the permute and the view as end edges - subgraph.edges_out.add((node, user)) + subgraph.edges_out.add((users_source, user)) # Mark the view for inclusion so it gets preserved continue return False @@ -420,6 +560,26 @@ def visit( # noqa: C901 return True + def _absorb_interleave( + self, + head: torch.fx.Node, + triple: tuple[int, int, torch.fx.Node, torch.fx.Node], + subgraph: Subgraph, + processed_nodes: set[torch.fx.Node], + current_end_permute: list[int], + current_start_permute: list[int], + ) -> torch.fx.Node | None: + """Add a matched interleave's interior nodes and return its tail.""" + _, _, expand_node, view_node = triple + if expand_node in processed_nodes or view_node in processed_nodes: + return None + for interior in (expand_node, view_node): + subgraph.nodes.add(interior) + subgraph.node_end_permute[interior] = current_end_permute + subgraph.node_start_permute[interior] = current_start_permute + subgraph.interleaves[head] = triple + return view_node + def _is_constant(self, node: torch.fx.Node) -> bool: """Check if a node's value is available at compile time. Only considers direct constants (get_attr, parameter/buffer/constant @@ -491,12 +651,29 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 if not self._subgraph_edges_are_current(subgraph): return False + # Nodes belonging to a repeat_interleave triple are rewritten as a unit + # below, so they must skip the per-node dim handling and the view rank + # check (the triple's interior ranks intentionally differ from the + # region's permutation rank). + interleave_nodes: set[torch.fx.Node] = set() + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): + interleave_nodes.update((head, expand_node, view_node)) + perm = subgraph.node_start_permute.get(head, subgraph.start_permute) + inp = head.args[0] if head.args else None + if not isinstance(inp, torch.fx.Node): + return False + in_shape = self._concrete_shape(inp) + if in_shape is None or len(perm) != len(in_shape): + return False + # Validate: every view_copy node's permutation rank must match its # input tensor rank. A mismatch can occur when a squeeze/unsqueeze # view is reached via upstream traversal with a permutation that was # already adapted to a different rank. Applying the optimisation in # this case would produce an invalid graph, so skip the subgraph. for node in subgraph.nodes: + if node in interleave_nodes: + continue if node.target in self._VIEW_OPS: perm = subgraph.node_start_permute.get(node, subgraph.start_permute) inp = node.args[0] @@ -507,6 +684,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Handle dimension related node arguments FIRST, before # bypassing permutes (which changes node inputs/metadata). for node in subgraph.nodes: + if node in interleave_nodes: + continue node_start_perm = subgraph.node_start_permute.get( node, subgraph.start_permute ) @@ -524,6 +703,13 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 elif node.target in self._VIEW_OPS: self.update_view_copy(node, node_start_perm) + for head, triple in subgraph.interleaves.items(): + self.update_interleave( + head, + triple, + subgraph.node_start_permute.get(head, subgraph.start_permute), + ) + # Skip incoming permutes. for inp, out in subgraph.edges_in: assert inp.target == exir_ops.edge.aten.permute_copy.default @@ -594,8 +780,51 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if const_node not in user_node.all_input_nodes: return False + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): + if ( + len(head.users) != 1 + or len(expand_node.users) != 1 + or expand_node not in head.users + or view_node not in expand_node.users + ): + return False + return True + def update_interleave( + self, + head: torch.fx.Node, + triple: tuple[int, int, torch.fx.Node, torch.fx.Node], + start_permute: list[int], + ) -> None: + """Retarget a repeat_interleave triple at the un-permuted layout. + + After the boundary permutes are removed the triple's input is in the + original layout, so the dim it interleaves moves from ``dim`` to + ``start_permute[dim]`` and all three shape arguments are rebuilt there. + """ + dim, scale, expand_node, view_node = triple + inp = cast(torch.fx.Node, head.args[0]) + in_shape = [int(d) for d in inp.meta["val"].shape] + inverse_permute = [start_permute.index(i) for i in range(len(start_permute))] + unpermuted_in = [in_shape[inverse_permute[i]] for i in range(len(in_shape))] + target_dim = start_permute[dim] + + if head.target == exir_ops.edge.aten.unsqueeze_copy.default: + set_arg(head, "dim", target_dim + 1) + else: + unsqueezed = list(unpermuted_in) + unsqueezed.insert(target_dim + 1, 1) + set_arg(head, "size", unsqueezed) + + expand_size = list(unpermuted_in) + expand_size.insert(target_dim + 1, scale) + set_arg(expand_node, "size", expand_size) + + merged = list(unpermuted_in) + merged[target_dim] *= scale + set_arg(view_node, "size", merged) + def update_cat(self, node: torch.fx.Node, start_permute: list[int]) -> None: dim = get_arg(node, "dim", int) set_arg(node, "dim", start_permute[dim]) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 491923e36a5..4f47912d4ed 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1492,6 +1492,335 @@ def test_broadcast_rank_increase_no_crash(self) -> None: # ────────────────────────────────────────────────────────────────────── +class RemovePermutesAroundRepeatInterleaveTest(unittest.TestCase): + """repeat_interleave lowers to unsqueeze -> expand_copy -> merging view_copy. + + The triple is rank-preserving, so a permutation flows through it once the + interleaved dim is remapped. Shapes here mirror torchaudio's Stretch2d as it + appears in wavernn between two channels-last convolutions. + """ + + @staticmethod + def _interleave( + builder: GraphBuilder, + inp: object, + shape: list[int], + dim: int, + scale: int, + ) -> object: + """Emit unsqueeze(dim+1) -> expand_copy(scale) -> view_copy(merge).""" + unsqueezed = list(shape) + unsqueezed.insert(dim + 1, 1) + expanded = list(unsqueezed) + expanded[dim + 1] = scale + merged = list(shape) + merged[dim] *= scale + + u = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(inp, dim + 1) + ) + e = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, args=(u, expanded) + ) + return builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(e, merged) + ) + + def test_removes_permutes_around_repeat_interleave(self) -> None: + """permute(NHWC->NCHW) -> interleave(W) -> permute(NCHW->NHWC): + both permutes should cancel and the interleave move to the NHWC dim.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + v = self._interleave(builder, p1, [1, 1, 16, 20], dim=3, scale=2) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.expand_copy.default), 1 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "RepeatInterleavePermutes" + ) + + def test_repeat_interleave_with_keyword_arguments(self) -> None: + """The matcher and rewrite support schema arguments passed by keyword.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + u = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, + args=(p1,), + kwargs={"dim": 4}, + ) + e = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, + args=(u,), + kwargs={"size": [1, 1, 16, 20, 2]}, + ) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(e,), + kwargs={"size": [1, 1, 16, 40]}, + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "RepeatInterleaveKwargs" + ) + + def test_repeat_interleave_rank_mismatch_is_not_rewritten(self) -> None: + """Reject a triple whose active boundary permutation has the wrong rank.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + interleaved = self._interleave(builder, p1, [1, 1, 16, 20], dim=3, scale=2) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(interleaved, [0, 2, 3, 1]), + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + start_permute = original.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + )[0] + + class RankMismatchPass(RemovePermutesAroundElementwiseOps): + def get_permutation(self, node: torch.fx.Node) -> list[int] | None: + if node is start_permute: + return [0, 2, 1] + return super().get_permutation(node) + + result = cast(PassResult, RankMismatchPass()(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_data], + "RepeatInterleaveRankMismatch", + ) + + def test_repeat_interleave_after_nop_stretch(self) -> None: + """The wavernn region verbatim: a freq_scale=1 Stretch2d leaves a nop + unsqueeze/view pair ahead of the real interleave. The permutation must + round-trip across it and still reach the triple.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + u0 = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 3) + ) + v0 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(u0, [1, 1, 16, 20]) + ) + v = self._interleave(builder, v0, [1, 1, 16, 20], dim=3, scale=2) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = _canonicalize_and_remove_permutes(original) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "RepeatInterleaveNopStretch" + ) + + def test_repeat_interleave_scale_four_on_middle_dim(self) -> None: + """Interleaving a non-trailing dim with a scale other than 2.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 3) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + v = self._interleave(builder, p1, [1, 3, 16, 20], dim=2, scale=4) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "RepeatInterleaveMiddleDim" + ) + + def test_repeat_interleave_composes_with_elementwise(self) -> None: + """An interleave and a pointwise op in the same permuted region.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(p1, p1)) + v = self._interleave(builder, mul, [1, 1, 16, 20], dim=3, scale=2) + relu = builder.call_operator(op=exir_ops.edge.aten.hardtanh.default, args=(v,)) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(relu, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "RepeatInterleaveElementwise" + ) + + def test_resnet_stretch_region(self) -> None: + """wavernn's resnet_stretch: the region enters through an unsqueeze at a + position the permutation moves, and leaves through squeeze_copy.dims.""" + builder = GraphBuilder() + x_data = torch.randn(1, 8, 16) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + u1 = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 1) + ) + # freq_scale=1 Stretch2d leaves a nop unsqueeze/view pair. + u2 = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(u1, 3) + ) + v2 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(u2, [1, 1, 16, 8]) + ) + v3 = self._interleave(builder, v2, [1, 1, 16, 8], dim=3, scale=4) + sq = builder.call_operator( + op=exir_ops.edge.aten.squeeze_copy.dims, args=(v3, [1]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(sq, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = _canonicalize_and_remove_permutes(original) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, result.graph_module, [x_data], "ResnetStretchRegion" + ) + + def test_non_merging_view_after_expand_is_not_optimized(self) -> None: + """A view that is not the (dim, dim+1) merge is not layout-invariant, + so the region must be left alone.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + u = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 4) + ) + e = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, args=(u, [1, 1, 16, 20, 2]) + ) + # Squeezes the leading dim instead of merging dims 3 and 4. + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(e, [1, 16, 20, 2]) + ) + builder.output([v]) + 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), 1 + ) + + def test_expand_with_extra_user_is_not_optimized(self) -> None: + """Rewriting the triple in place would corrupt a second consumer of the + expand, so the triple must not be claimed.""" + builder = GraphBuilder() + x_data = torch.randn(1, 16, 20, 1) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + u = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 4) + ) + e = builder.call_operator( + op=exir_ops.edge.aten.expand_copy.default, args=(u, [1, 1, 16, 20, 2]) + ) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(e, [1, 1, 16, 40]) + ) + other = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(e, e)) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2, other]) + 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), 2 + ) + + class RemovePermutesAroundElementwiseOpsTest(unittest.TestCase): def test_no_permutes_is_noop(self) -> None: """With no surrounding permutes, the pass makes no change."""