Recognise the repeat_interleave idiom in RemovePermutesAroundElementwiseOps (pytorch#21790)

Differential Revision: D114508262

Pull Request resolved: pytorch#21790

Author: mcremon-meta

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,19 +482,19 @@ 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.
                     user_users = list(user.users.keys())
                     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])