Skip to content

Fold non-matching end permutes instead of discarding the region (#22293) #22293

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

Merged
meta-codesync merged 1 commit into from
Sep 1, 2026
+251 −74
Merged

Fold non-matching end permutes instead of discarding the region (#22293) #22293

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
14 changes: 12 additions & 2 deletions backends/cadence/aot/tests/test_remove_ops_passes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,11 +477,21 @@ def test_keep_permutes_around_elemwise_ops_add(self) -> None:
)
builder.output([permute])
original = builder.get_graph_module()
gm_before = copy.deepcopy(original)
p = RemovePermutesAroundElementwiseOps()
graph_after_passes = cast(PassResult, p(original)).graph_module
# Ensure no permutes were removed, since the dimensions don't fit the expected pattern
# The end permute is not the inverse of the start one, so it cannot be
# dropped. It absorbs the start permute instead, leaving one behind.
self.assertEqual(
count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 2
count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 1
)

sample_inputs = [torch.randn(1, 8, 4, 4, dtype=torch.float32)]
validate(
gm_before,
graph_after_passes,
sample_inputs,
"RemovePermutesAroundElementwiseOps",
)

def test_remove_permutes_around_elemwise_ops_add_mean(self) -> None:
Expand Down
114 changes: 47 additions & 67 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 @@ -41,6 +41,12 @@ class Subgraph:
edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set)
# Outgoing edges of the subgraph to permute nodes.
edges_out: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set)
# Outgoing edges to permutes that do not match end_permute. Those are
# kept and their permutation rewritten to absorb the removed start
# permute, as (producer, permute node, new permutation).
edges_out_to_update: set[
tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]
] = field(default_factory=set)
# Incoming edges from constant nodes that need a compensating permute.
constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field(
default_factory=set
Expand Down Expand Up @@ -372,7 +378,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
# Expected end permutation for the subgraph.
end_permute = [start_permute.index(i) for i in range(len(start_permute))]

# Try direct users first (same-rank matching)
for user in node.users:
if (
not self.is_node_permutable(user)
Expand All @@ -385,50 +390,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
for n in subgraph.nodes:
processed_nodes.add(n)

# Also try: permute → view(squeeze/unsqueeze) → chain → ...
# If the permute's sole user is a squeeze/unsqueeze view,
# adapt the permutation across the view and search for a
# matching end permute at the new rank.
users = list(node.users.keys())
if (
len(users) == 1
and self._is_squeeze_unsqueeze_view(users[0])
and node not in processed_nodes
):
view_node = users[0]
adapted_start = self._adapt_permute_across_view(
start_permute, view_node
)
if adapted_start is not None:
adapted_end = [
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)
and self._interleave_triple(view_user) is None
):
continue
subgraph = self.Subgraph(adapted_start, adapted_end)
# Include the view in the subgraph
subgraph.nodes.add(view_node)
subgraph.node_end_permute[view_node] = adapted_end
# Use the ORIGINAL start_permute for the view node
# so update_view_copy can remap its shape correctly
subgraph.node_start_permute[view_node] = start_permute
# The start permute feeds into the view
subgraph.edges_in.add((node, view_node))
if self.visit(
view_user,
subgraph,
processed_nodes,
adapted_end,
adapted_start,
):
subgraphs_found.append(subgraph)
for n in subgraph.nodes:
processed_nodes.add(n)

modified = False
for subgraph in subgraphs_found:
if self.permute_subgraph(subgraph):
Expand Down Expand Up @@ -526,29 +487,17 @@ def visit( # noqa: C901
if user_perm == downstream_end:
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: 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(
downstream_start, view_after
# Non-matching permute: keep it and fold the start permute into it
# rather than discarding the region.
if user_perm is None or len(user_perm) != len(downstream_start):
return False
subgraph.edges_out_to_update.add(
(
users_source,
user,
tuple(downstream_start[d] for d in user_perm),
)
if adapted_start_after is not None:
adapted = [
adapted_start_after.index(i)
for i in range(len(adapted_start_after))
]
if user_perm == adapted:
# Include both the permute and the view as end edges
subgraph.edges_out.add((users_source, user))
# Mark the view for inclusion so it gets preserved
continue
return False
)
elif user.op == "output":
return False
elif self._is_permutation_sink_view(user):
Expand Down Expand Up @@ -682,11 +631,27 @@ def _is_constant_pad(self, node: torch.fx.Node) -> bool:

return True

def _removes_a_permute(self, subgraph: Subgraph) -> bool:
"""Whether rewriting this region reduces the number of permutes."""
if subgraph.edges_out:
return True
rewired = set(subgraph.edges_in)
for permute, _ in subgraph.edges_in:
if all((permute, user) in rewired for user in permute.users):
return True
return False

def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901
# Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes.
if not self._subgraph_edges_are_current(subgraph):
return False

# Folding an end permute only pays for itself if some permute goes away.
# Otherwise the region is rewritten for nothing, and the composed
# permutation is a worse fusion candidate for the passes downstream.
if subgraph.edges_out_to_update and not self._removes_a_permute(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
Expand Down Expand Up @@ -794,6 +759,11 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901
assert out.target in PERMUTE_COPY_TARGETS
out.replace_all_uses_with(inp)

# Update outgoing permutes that can't be eliminated.
for _, out, new_permutation in subgraph.edges_out_to_update:
assert out.target in PERMUTE_COPY_TARGETS
set_arg(out, "dims", list(new_permutation))

return True

def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
Expand All @@ -802,10 +772,20 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes:
return False

# edges_out_to_update can rewrite a permute in place, leaving it wired.
if self.get_permutation(inp) != subgraph.node_start_permute.get(
out, subgraph.start_permute
):
return False

for inp, out in subgraph.edges_out:
if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users:
return False

for inp, out, _ in subgraph.edges_out_to_update:
if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users:
return False

for const_node, user_node in subgraph.constant_edges_in:
if const_node not in user_node.all_input_nodes:
return False
Expand Down
Loading
Loading