[None][refactor] AutoDeploy: make MLIR elementwise fusion pipeline xDSL-native - #15130
Open
suyoggupta wants to merge 1 commit into
Open
[None][refactor] AutoDeploy: make MLIR elementwise fusion pipeline xDSL-native#15130suyoggupta wants to merge 1 commit into
suyoggupta wants to merge 1 commit into
Conversation
…SL-native Addresses PR NVIDIA#12427 review (comment r3019035416): the fusion pipeline was half-idiomatic — decompose.py used xDSL's PatternRewriteWalker but the fusion half hand-rolled IR traversal, mutation, and cleanup, including block.erase_op(safe_erase=False) which suppresses use-def validation. Changes: - subgraph_replace.py: remove safe_erase=False. Subgraph ops are erased in reverse-topological order (consumers before producers) after replace_by redirects external uses, so the default safe_erase=True check passes. Added an optional PatternRewriter param: the walker path routes insertion/erasure through the rewriter; standalone callers (unit tests) keep using block ops. - New mlir/fusion/fuse.py: run_fusion() + FusionStats, mirroring run_decomposition. A _FuseSubgraphPattern (RewritePattern) matches each subgraph's anchor and is driven by PatternRewriteWalker(apply_recursively= False). The transform _apply now just calls run_fusion. Discovery deliberately stays the custom greedy union-find (plus placement-conflict detection, the 64-input split, and Kahn topo-sort for hash stability): a root-anchored backward walk would fragment maximal multi-output subgraphs and cannot express the dependent 64-input partitioning. Only iteration and mutation become xDSL-native. A forward walk preserves partition dependency order so refresh_inputs() still works. Tests: 3 new tests in test_elementwise_fusion_e2e.py exercise the run_fusion walker path (structural/metadata, low-rank skip, CUDA full FX roundtrip). Full MLIR suite: 46 passed. Also documents the change in mlir/agent_learnings.md (#15) and adds an ad-mlir-fusion-update skill distilling that log into an update/debug guide. Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR refactors the MLIR elementwise fusion pass to use xDSL's ChangesWalker-driven fusion pass refactoring
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py (1)
506-626: Coverage is sufficient for the walker-driven fusion refactor scope.
tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.pynow covers:
- walker replacement + metadata contract,
- low-rank skip accounting,
- CUDA FX→MLIR→run_fusion→FX numerical roundtrip.
No must-have follow-up test file is required for this PR layer.
As per coding guidelines, test reviews under
tests/**should explicitly assess whether coverage is sufficient/insufficient for the changed behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py` around lines 506 - 626, The test file lacks an explicit coverage assessment statement required by our guidelines; update tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py to include a short explicit coverage assertion/comment indicating these three tests (test_run_fusion_replaces_subgraph_and_records_metadata, test_run_fusion_skips_low_rank_subgraph, test_run_fusion_full_fx_roundtrip) collectively cover walker-driven fusion replacement, low-rank skipping, and CUDA FX→MLIR→run_fusion→FX numerical roundtrip — add this as a one-line comment near the top of the test file or immediately above the first of these test functions so reviewers can see the coverage assessment.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.claude/skills/ad-mlir-fusion-update/SKILL.md:
- Around line 146-150: The example command doesn't explicitly enable the fusion
pass so it won't run with default settings; update the example invocation of
build_and_run_ad.py (the script name build_and_run_ad.py and the existing
--args.model flag) to include an explicit config overlay or flag that turns the
fusion transform on (for example append an explicit config-overlay or
--enable-fusion/--args.overlay JSON that sets the fusion pass to true) so the
docs demonstrate an end-to-end run that actually exercises the fusion pass.
---
Nitpick comments:
In `@tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py`:
- Around line 506-626: The test file lacks an explicit coverage assessment
statement required by our guidelines; update
tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py to
include a short explicit coverage assertion/comment indicating these three tests
(test_run_fusion_replaces_subgraph_and_records_metadata,
test_run_fusion_skips_low_rank_subgraph, test_run_fusion_full_fx_roundtrip)
collectively cover walker-driven fusion replacement, low-rank skipping, and CUDA
FX→MLIR→run_fusion→FX numerical roundtrip — add this as a one-line comment near
the top of the test file or immediately above the first of these test functions
so reviewers can see the coverage assessment.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ab6a29b0-2bb1-4528-bee4-e0fec94eef19
📒 Files selected for processing (6)
.claude/skills/ad-mlir-fusion-update/SKILL.mdtensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.mdtensorrt_llm/_torch/auto_deploy/mlir/fusion/fuse.pytensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.pytensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.pytests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py
Comment on lines
+146
to
+150
| Then end-to-end (the pass is off by default — enable it or use a config overlay): | ||
|
|
||
| ```bash | ||
| python examples/auto_deploy/build_and_run_ad.py --args.model=<HF_MODEL> | ||
| ``` |
Contributor
There was a problem hiding this comment.
Validation command should explicitly enable the fusion pass.
This command currently won’t exercise fusion when defaults are unchanged, because the same doc states the transform is disabled by default. Please include an explicit config overlay/flag in the example command.
Suggested doc tweak
- python examples/auto_deploy/build_and_run_ad.py --args.model=<HF_MODEL>
+ python examples/auto_deploy/build_and_run_ad.py \
+ --args.model=<HF_MODEL> \
+ --args.yaml-extra=<CONFIG_WITH_MLIR_ELEMENTWISE_FUSION_ENABLED>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In @.claude/skills/ad-mlir-fusion-update/SKILL.md around lines 146 - 150, The
example command doesn't explicitly enable the fusion pass so it won't run with
default settings; update the example invocation of build_and_run_ad.py (the
script name build_and_run_ad.py and the existing --args.model flag) to include
an explicit config overlay or flag that turns the fusion transform on (for
example append an explicit config-overlay or --enable-fusion/--args.overlay JSON
that sets the fusion pass to true) so the docs demonstrate an end-to-end run
that actually exercises the fusion pass.
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #53169 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #53169 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Addresses the review on PR #12427 (comment r3019035416): the MLIR elementwise fusion pipeline was half-idiomatic.
mlir/decompose.pyalready used xDSL'sPatternRewriteWalker+GreedyRewritePatternApplier+RewritePattern, but the fusion half (subgraph_discovery.py+subgraph_replace.py) hand-rolled IR traversal, mutation, and cleanup — includingblock.erase_op(..., safe_erase=False), which suppresses use-def validation.What changed
mlir/fusion/subgraph_replace.py: removedsafe_erase=False. Subgraph ops are erased in reverse-topological order (consumers before producers); combined with thereplace_bythat redirects every external use, each op has no remaining uses at erase time, so the defaultsafe_erase=Trueuse-def check passes. Added an optionalrewriter: PatternRewriterparam — the walker path routes insertion/erasure through the rewriter; standalone callers (unit tests) keep using block ops.mlir/fusion/fuse.py(new):run_fusion()+FusionStats, mirroringrun_decomposition. A_FuseSubgraphPattern(RewritePattern) matches each subgraph's anchor (sg.ops[0]) and is driven byPatternRewriteWalker(apply_recursively=False). The transform_applynow simply callsrun_fusion.transform/library/mlir_elementwise_fusion.py:_applyslimmed — the manualfor sg in subgraphsloop collapses into onerun_fusion(...)call. Top-level pipeline now reads symmetrically:run_decomposition→run_fusion.Why discovery stays custom (not the sketched root-walk)
Subgraph discovery deliberately remains the greedy union-find (plus
_has_placement_conflict, the 64-input_split_subgraph, and Kahn topo-sort for hash stability). A root-anchored backward walk would (1) fragment maximal multi-output connected components into separate fusions (different hashes, more kernels) and (2) cannot express the dependent 64-input partitioning. So only iteration + mutation became xDSL-native. A forward walk preserves split-partition dependency order, sorefresh_inputs()continues to work exactly as before.Test Coverage
tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.pyexercise therun_fusionwalker path: structural/metadata, low-rank skip, and a CUDA full FX roundtrip with numerical check.Docs
mlir/agent_learnings.md: new entry Fix the link to the documentation #15 documenting the port and the keep-discovery-custom rationale.ad-mlir-fusion-updateClaude Code skill distilling the learnings log into an update/debug guide for the pass.🤖 Generated with Claude Code
Summary by CodeRabbit
Documentation
New Features
Tests