Skip to content

Recognise the repeat_interleave idiom in RemovePermutesAroundElementwiseOps (#21790) #21790

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mcremon-meta wants to merge 1 commit into
base: main
Choose a base branch
from
+577 −19
Open

Recognise the repeat_interleave idiom in RemovePermutesAroundElementwiseOps (#21790) #21790

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
267 changes: 248 additions & 19 deletions backends/transforms/remove_permutes_around_elementwise_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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__()
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand All @@ -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
)
Expand All @@ -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
Expand Down Expand Up @@ -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])
Expand Down
Loading
Loading