Add CSEPass to the default edge compilation pipeline (#19849)

[billmguo]

Summary:

CSEPass (Common Subexpression Elimination) was previously only used by specific backends (MLX, Cadence) and export scripts. This change adds it to base_post_op_replace_passes in the default to_edge / to_edge_transform_and_lower pipeline so all backends benefit from CSE automatically.

The pass is placed before dead_code_elimination_pass so any dead nodes CSE creates get cleaned up, and before DebugHandleGeneratorPass so debug handles are assigned to the final deduplicated graph.

CSE can merge structurally-identical nodes across different source partitions, which caused ConvertToLinearPass in XNNPACK to see extra output nodes in a partition. This is fixed by filtering outputs to only those reachable from the mm/addmm node within the partition.

The test_dim_order_revert_pass test is updated to expect 2 _to_copy ops (instead of 4) after to_edge, since CSE deduplicates the two identical branches.

Reviewed By: mohankumarkumar

Differential Revision: D106584552

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19849

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 1 Unrelated Failure

As of commit f1614d2 with merge base 4469d84 ():

NEW FAILURES - The following jobs have failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 18397797c17f4431b410e74707753509bdde3d49526ce87b924d750028fd3f1b /exec failed with exit code 1
• trunk / test-arm-backend-ethos-u (test_pytest_ops_ethos_u85) / linux-job (gh)
  RuntimeError: Command docker exec -t c72ff089db4e735fb70f3b374f59ac338d1ec7febdb9650deb389b78ffb3d063 /exec failed with exit code 1
• trunk / test-arm-backend-vkml (test_pytest_ops_vkml) / linux-job (gh)
  RuntimeError: Command docker exec -t 7b77e3aa53352da6e2b37f8fef6564f90acf140b5621781230448133f83015b6 /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• trunk / test-arm-backend-ethos-u (test_smaller_stories_llama) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

• ❌ - login: @billmguo / name: Min Guo. The commit (f1614d2) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[meta-codesync]

@billmguo has exported this pull request. If you are a Meta employee, you can view the originating Diff in D106584552.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@metascroy can you take a look?

@claude review this PR

@billmguo Can you please sign the CLA?

[claude]

Claude finished @digantdesai's task in 1m 37s —— View job

---

Review of PR #19849: Add CSEPass to the default edge compilation pipeline

• Gather context and read changed files
• Review exir/passes/__init__.py changes
• Review backends/xnnpack/_passes/convert_to_linear.py changes
• Review exir/passes/BUCK changes
• Review test changes
• Post review feedback

---

Overall: This is a well-structured change. Adding CSE to the default pipeline is a good idea — deduplicating identical computations benefits all backends without requiring them to opt in individually. The placement before DCE and DebugHandleGeneratorPass is correct. The XNNPACK fix is well-motivated. A few observations below:

---

backends/xnnpack/_passes/convert_to_linear.py — Reachability fix

The forward-reachability walk from node.users (lines 132–138) correctly filters output nodes to only those downstream of the mm/addmm node within the partition. However:

1. BFS vs DFS semantics: Using queue.pop() gives DFS behavior (stack). This is fine for correctness since you just need reachability, but naming the variable queue is slightly misleading. Minor nit — no functional issue.

2. Potential infinite loop safety: The if cur in reachable guard correctly prevents revisiting nodes. Good.

3. Edge case — zero reachable outputs: If CSE merges nodes such that no output in src_partition.output_nodes is reachable from the mm/addmm, the subsequent assert will fire with an empty list. This seems unlikely given how source partitions work, but worth noting that the assert message would show [] which is still informative. Fine as-is.

backends/arm/test/passes/test_fuse_duplicate_users_pass.py

The ops_before_pass values now equal ops_after_pass for both FuseaAvgPool (1 → 1) and FuseAvgPoolChain (2 → 2). This means the FuseDuplicateUsersPass is effectively a no-op in these tests now that CSE runs first during to_edge. The tests still validate that the pass doesn't break anything when there's nothing to fuse, but they no longer test the pass's fusing behavior.

Suggestion: Consider adding a test case where inputs differ (so CSE can't deduplicate) but outputs are consumed by the same user — that would still exercise the fuse pass's core logic independently of CSE.

exir/passes/__init__.py

Placement looks correct:

CSEPass()           # deduplicate identical subexpressions
dead_code_elimination_pass  # clean up dead nodes CSE may leave behind
DebugHandleGeneratorPass()  # assign handles to final graph

CSE already calls graph.eliminate_dead_code() internally when it modifies the graph (line 122 of cse_pass.py), so the subsequent dead_code_elimination_pass is redundant when CSE fires — but harmless and ensures cleanup if other passes between pre/post introduce dead code. Fine.

exir/tests/test_passes.py

The test update for test_dim_order_revert_pass is well-commented. The observation that CSE doesn't deduplicate dim_order ops (because they're non-aten and treated as unsafe by _is_safe_target) is correct and important — it means the _skip_dim_order=False path remains unaffected by CSE.

CSEPass implementation (exir/passes/cse_pass.py)

The implementation is solid:

• Value numbering approach is clean and O(n) per node
• Schema introspection for safety is principled — only trusting aten:: schemas
• The _output_nodes protection prevents deduplicating buffer mutation outputs (important for KV cache patterns)
• The TypeError catch in _make_hashable is a good fallback for exotic arg types

One minor observation: _is_safe caches by id(target) (line 129), which is correct since op targets are interned singletons in PyTorch. Just calling it out as an assumption that holds today.

Summary

Approve with minor suggestion: The change is correct and well-thought-out. The only actionable suggestion is to consider whether the ARM FuseDuplicateUsersPass tests still provide meaningful coverage now that CSE pre-empts their fusing — but that's arguably a follow-up concern, not a blocker.

---