diff --git a/.claude/skills/ad-mlir-fusion-update/SKILL.md b/.claude/skills/ad-mlir-fusion-update/SKILL.md new file mode 100644 index 000000000000..9468e05a8d5c --- /dev/null +++ b/.claude/skills/ad-mlir-fusion-update/SKILL.md @@ -0,0 +1,175 @@ +--- +name: ad-mlir-fusion-update +description: > + Modify, extend, or debug the TensorRT-LLM AutoDeploy MLIR elementwise fusion pass + (tensorrt_llm/_torch/auto_deploy/mlir/) — the FX→MLIR→decompose→discover→Triton-codegen→replace→FX + pipeline built on xDSL. Use when adding a dialect op/primitive, editing the Triton emitter, changing + subgraph discovery/splitting, the FX↔MLIR converters, or the rewrite/replacement plumbing; or when + triaging fusion failures (KeyError/ValueError in codegen or back-conversion, illegal memory access in + generated kernels, lost GraphModule methods). Covers the 4-touchpoint rule, fail-safe skip, a + failure-mode triage table, the validation suite, and the agent_learnings.md log convention. +license: Apache-2.0 +tags: + - tensorrt-llm + - autodeploy + - mlir + - xdsl + - fusion + - triton +metadata: + author: NVIDIA Corporation +--- + +# AutoDeploy: Update the MLIR Elementwise Fusion Pass + +## When to use this skill + +Use this when the task touches the **MLIR elementwise fusion subsystem** under +`tensorrt_llm/_torch/auto_deploy/mlir/` — adding a primitive op, changing the +Triton emitter, subgraph discovery/splitting, the FX↔MLIR converters, the +rewrite plumbing, or triaging a fusion crash. + +For adding a *generic* fusion transform under `transform/library/` (not the +MLIR pipeline), use **ad-add-fusion-transformation** instead. For dumping +before/after graphs, use **ad-graph-dump**. For checking whether the pass +actually ran in a serve log, use **ad-conf-check**. + +> **Canonical log:** `tensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.md` +> is the chronological issue/fix record (currently 15 entries). **Read it first** +> — this skill distills it. **After fixing a new bug, append a numbered entry** +> (Symptom / Root cause / Fix / Lesson) so the log stays the source of truth. + +## Architecture map + +The transform `mlir_elementwise_fusion` (in +`transform/library/mlir_elementwise_fusion.py`, registry key +`mlir_elementwise_fusion`) runs a 6-step pipeline. It is **fail-safe**: any +unsupported dtype/op/pattern must cause a graceful skip, never a crash. + +| Step | Code | Notes | +|------|------|-------| +| 1. FX → MLIR | `mlir/fx_to_mlir.py` (`FXToMLIRConverter`) | Builds the `ad` dialect IR + a metadata side-table. Unmodeled ops → `ad.opaque`. | +| 2. Decompose | `mlir/decompose.py` (`run_decomposition`) + `mlir/decompose_rules/` | RMSNorm/gated-RMSNorm → elementwise primitives. xDSL `PatternRewriteWalker`. | +| 3+4. Discover + replace | `mlir/fusion/fuse.py` (`run_fusion`) → `subgraph_discovery.py` + `subgraph_replace.py` | Walker-driven. Discovery is custom union-find; replacement is rewriter-tracked. | +| (codegen) | `mlir/codegen/triton_emitter.py` (`_EMIT` table) + `kernel_cache.py` | Generates + caches one Triton kernel per subgraph. | +| 6. MLIR → FX | `mlir/mlir_to_fx.py` (`MLIRToFXConverter`) | Rebuilds the GM with fused kernel calls. | +| dialect | `mlir/dialect.py` (`AD_OPS`, dtype maps) | Op + type definitions; `_TORCH_TO_MLIR_DTYPE`. | + +`default.yaml` entry: `mlir_elementwise_fusion` (stage `post_load_fusion`, +`enabled: false` by default, `run_shape_prop: false`). xDSL is optional — +everything is gated behind `HAS_XDSL`; the pass skips cleanly when absent. + +## The 4-touchpoint rule (adding a primitive op) + +When you add a new elementwise/reduction primitive, update **all four** or +back-conversion will silently drop it (agent_learnings #1): + +1. `mlir/dialect.py` — define the `AdXxx` op class; add to `AD_OPS`. Add any new + dtype to `_TORCH_TO_MLIR_DTYPE` (and the reverse map). +2. `mlir/fx_to_mlir.py` — route the aten op in `_convert_call_function`. +3. `mlir/mlir_to_fx.py` — add the reverse handler. +4. `mlir/codegen/triton_emitter.py` — add an `_EMIT` entry (or a special-case + branch for attribute-carrying ops like `pow`/`splat`/`cast`/`floordiv`/`eq`). + To make it fusible, also add the class to `FUSIBLE_OPS` in + `subgraph_discovery.py`. + +## Cardinal invariants (do not regress) + +These are the load-bearing rules behind the log — violating one reintroduces a +known crash. + +- **Fail-safe over correctness-at-all-costs.** Unsupported dtype/op/pattern → + log a warning and skip (return original graph). The `converter.convert()` call + is wrapped in `try/except ValueError` for exactly this (#13). +- **Discovery stays union-find — do NOT replace it with a root-anchored walk** + (#15). A backward root cone fragments maximal multi-output components and + can't express the 64-input dependent partitioning. Only iteration + mutation + are xDSL-native; the discovery algorithm (union-find, `_has_placement_conflict`, + `_split_subgraph`, Kahn topo-sort) is custom by design. +- **Safe erase via reverse-topological order** (#15). `subgraph.ops` is topo-sorted; + erase in reverse (consumers first) after `replace_by` redirects external uses — + then default `safe_erase=True` passes. Never reach for `safe_erase=False`. +- **Forward walk only.** `run_fusion` uses `PatternRewriteWalker(apply_recursively=False)` + with a **forward** walk so split-partition anchors are visited in dependency + order and `refresh_inputs()` (#10) still works. Do not set `walk_reverse=True`. +- **`run_shape_prop: false`** for this transform (#6) — shape-prop concretizes + symbolic dims and breaks downstream transforms. Fused `register_fake` handles + shapes via the reference-input mechanism. +- **Three input load categories in the emitter** (#3, #4): broadcast (lower rank → + `offs` only), narrow (same rank, last-dim < `N_COLS` → scalar load), full-row + (last-dim == `N_COLS` → `row_off + offs`). `N_COLS` = **widest** highest-rank input. +- **Separate load vs store offsets** (#5): loads use `row_off` (from input + `stride(-2)`, may be non-contiguous from `aten.chunk`); stores use + `out_row_off = pid * N_COLS` (outputs are always contiguous). +- **64-input limit is hard** (#9): `torch.library.custom_op` schemas cap at 64 + args. `_split_subgraph` must split oversized subgraphs **placement-aware** (#11): + a partition is viable only if all input producers precede all output consumers. +- **`replace_subgraph_with_fused_op`'s `rewriter` param must stay optional** (#15) — + unit tests call it standalone without a walker. +- **MLIR→FX must preserve GraphModule methods** (#14): copy missing public + callables (e.g. `get_input_embeddings`) onto the reconstructed GM. +- **Pattern-matchers downstream are fragile** (#7): fusion absorbs `aten.add.Tensor` + etc.; transforms like `multi_stream_moe` must match op-type-agnostically. + +## Failure-mode triage table + +| Symptom | Likely cause | See | +|---------|--------------|-----| +| `ValueError: FX node not found for MLIR value` (unfused op) | New op missing from a converter — 4-touchpoint violation | #1 | +| `ValueError: FX node not found ... 'ad.opaque' (all_reduce/dist)` | Fused op placed before an input producer in a spanning subgraph | #2, #11 | +| `RuntimeError: shape '[...]' is invalid` after fusion | `_get_ncols` not using the widest input | #3 | +| `Triton ... illegal memory access` | Narrow input loaded as full-row, or load/store stride mismatch | #4, #5 | +| `KeyError` in `triton_emitter` (`val_names` missing operand) | Stale subgraph inputs (missing `refresh_inputs`) or bad intra-subgraph topo order | #10, #12 | +| `RuntimeError: schema has N arguments but ... supports 64` | Oversized subgraph not split | #9 | +| Downstream sees concrete shapes instead of symbolic | `run_shape_prop: true` | #6 | +| `ValueError: Unsupported torch dtype ... complex64` | Unmapped dtype; should skip gracefully not crash | #13 | +| `AttributeError: 'GraphModule' has no attribute '...'` | Methods dropped on MLIR→FX reconstruction | #14 | +| `multi_stream_moe` produces 0 matches when fusion on | Merge-add absorbed by fusion; matcher too specific | #7 | + +## Validation + +Run the dedicated unit suite (xDSL required; CUDA tests auto-skip without a GPU): + +```bash +python3 -m pytest tests/unittest/auto_deploy/singlegpu/mlir/ -q +``` + +What the suite covers (add to the matching file when you change behavior): + +- `test_dialect.py` — op/type definitions. +- `test_fx_mlir_roundtrip.py` — FX↔MLIR conversion fidelity. +- `test_decompose.py` — decomposition rules. +- `test_subgraph_discovery.py` — grouping, splitting, **subgraph count** (parity invariant). +- `test_string_codegen.py` — generated kernel source + **hash stability** (kernel cache). +- `test_elementwise_fusion_e2e.py` — full pipeline incl. `run_fusion` walker path and CUDA numerical roundtrips. + +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= +``` + +Verify in logs: `Decomposed N high-level ops`, `Discovered N fusible subgraphs`, +`replaced M (skipped K low-rank)`, no `RuntimeError`/`KeyError`, correct output. +Always benchmark e2e **under CUDA graphs**, not standalone kernel speed (#8). + +After Python edits: `pre-commit run --files `. + +## Review checklist + +- New primitive? All **4 touchpoints** updated and added to `FUSIBLE_OPS` if fusible. +- No `safe_erase=False`; erase order is reverse-topological. +- Discovery algorithm unchanged unless intentionally so (no root-walk substitution). +- `run_fusion` stays forward-walk; `rewriter` param on `replace_subgraph_with_fused_op` still optional. +- Unsupported input → graceful skip (covered by a test), not a crash. +- Subgraph-count and hash-stability tests still pass (no kernel-cache churn). +- New bug fixed? **Appended a numbered entry to `agent_learnings.md`.** + +## Guardrails + +- One logical change per patch; don't mix a converter fix with an emitter rewrite. +- If the symptom isn't in the triage table, reproduce with `AD_DUMP_GRAPHS_DIR` + (see **ad-graph-dump**) before hypothesizing — the failure is usually a + missing converter handler or a stale-input/placement issue, not "slow". +- Treat the pass as optional infrastructure: never let a fusion bug degrade a + model that would otherwise run — prefer skip. diff --git a/tensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.md b/tensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.md index 99ad2f675d9b..52271a635ef4 100644 --- a/tensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.md +++ b/tensorrt_llm/_torch/auto_deploy/mlir/agent_learnings.md @@ -388,3 +388,58 @@ tried first but broke DeepSeek-V3 due to lost shape/export metadata. reconstructed `GraphModule` must preserve not just the graph and parameters but also any callable methods that downstream code relies on. Test with multimodal models that have extra methods on their graph modules. + +______________________________________________________________________ + +## 15. Porting fusion plumbing to xDSL's PatternRewriteWalker (keep discovery custom) + +**Context:** PR #12427 review (comment r3019035416) flagged that the pipeline +was half-idiomatic: `decompose.py` used xDSL's `PatternRewriteWalker` + +`GreedyRewritePatternApplier` + `RewritePattern`, but `subgraph_discovery.py` + +`subgraph_replace.py` hand-rolled IR traversal, mutation, and cleanup — +including `block.erase_op(..., safe_erase=False)`, which suppresses use-def +validation. + +**What was done:** + +- Removed `safe_erase=False`. Because `subgraph.ops` is topologically sorted + and we erase in **reverse** order (consumers before producers), and every + external use was already redirected by `replace_by`, each op has no remaining + uses at erase time — so the default `safe_erase=True` check passes. The old + flag was defensive, not load-bearing. +- Added `mlir/fusion/fuse.py`: `run_fusion(module, metadata, log_warning)` + + `FusionStats`, mirroring `run_decomposition`. A `_FuseSubgraphPattern` + (`RewritePattern`) matches each subgraph's anchor (`sg.ops[0]`) and is driven + by `PatternRewriteWalker(apply_recursively=False)`. The transform `_apply` + now just calls `run_fusion`. +- `replace_subgraph_with_fused_op` gained an optional `rewriter: PatternRewriter` + param: walker path uses `rewriter.insert_op`/`erase_op`; standalone callers + (unit tests call it directly) keep using block ops. The param **must** stay + optional — `test_elementwise_fusion_e2e.py` invokes it without a walker. + +**Decision — discovery stays union-find, NOT a root-anchored walk.** The review +sketched a `discover_subgraph_rooted_at` backward walk. That is *not* a faithful +substitute: a single-root backward cone (1) fragments maximal **multi-output** +connected components (two fusible results feeding different external consumers +land in separate fusions → different hashes, more kernels), and (2) cannot +express the 64-input **dependent** partitioning from issues #9/#11 (a single +`match_and_rewrite` emitting one `AdOpaque` per root can't chain partitions). +So only **iteration + mutation** went xDSL-native; the domain logic +(union-find, placement-conflict detection `_has_placement_conflict`, the +64-input `_split_subgraph`, Kahn topo-sort for hash stability) stays custom — +matching the reviewer's own "what stays custom" carve-out. + +**Subtlety:** A **forward** walk visits partition anchors in dependency order +(partition 1 before partition 2), so `refresh_inputs()` (issue #10) still picks +up redirected operands exactly as the old ahead-of-time loop did. Don't switch +to `walk_reverse=True` — it would break partition input refresh. + +**API gotcha:** xDSL 0.66 `Block` has **no** `insert_op_at_location`. Use +`block.insert_op_before/after` (standalone) or +`rewriter.insert_op(op, InsertPoint.before/after(anchor))` (walker). + +**Lesson:** When making a hand-rolled pass xDSL-idiomatic, port the *plumbing* +(walker-driven iteration, rewriter-tracked mutation, safe erase) but keep the +codegen-constraint-specific *discovery algorithm* custom. Reverse-topological +erase + prior `replace_by` is sufficient for safe erase — no need to bypass the +use-def check. diff --git a/tensorrt_llm/_torch/auto_deploy/mlir/fusion/fuse.py b/tensorrt_llm/_torch/auto_deploy/mlir/fusion/fuse.py new file mode 100644 index 000000000000..485f48c9c80c --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/mlir/fusion/fuse.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Walker-driven elementwise fusion pass. + +Mirrors :func:`..decompose.run_decomposition`: discovery computes the maximal +fusible subgraphs (greedy union-find + codegen-constraint partitioning — see +:mod:`.subgraph_discovery`), and an xDSL ``RewritePattern`` driven by +``PatternRewriteWalker`` performs the per-subgraph codegen + replacement. This +keeps the top-level pipeline consistent with decomposition (both go through +xDSL's pass infrastructure) while leaving the domain-specific discovery +algorithm — placement-conflict detection, the 64-input split, and hash-stable +topological ordering — untouched. +""" + +import logging +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional + +from xdsl.dialects.builtin import ModuleOp, TensorType +from xdsl.ir import Operation +from xdsl.pattern_rewriter import ( + GreedyRewritePatternApplier, + PatternRewriter, + PatternRewriteWalker, + RewritePattern, +) + +from ..codegen.kernel_cache import KernelCache +from ..codegen.triton_emitter import generate_kernel_from_subgraph +from .subgraph_discovery import FusibleSubgraph, discover_fusible_subgraphs +from .subgraph_replace import replace_subgraph_with_fused_op + +logger = logging.getLogger(__name__) + + +@dataclass +class FusionStats: + """Summary of a :func:`run_fusion` invocation.""" + + num_subgraphs: int = 0 + num_replaced: int = 0 + num_skipped_low_rank: int = 0 + + +def _max_input_rank(sg: FusibleSubgraph) -> int: + return max( + (len(inp.type.get_shape()) for inp in sg.inputs if isinstance(inp.type, TensorType)), + default=0, + ) + + +def _min_output_rank(sg: FusibleSubgraph) -> int: + return min( + (len(out.type.get_shape()) for out in sg.outputs if isinstance(out.type, TensorType)), + default=0, + ) + + +class _FuseSubgraphPattern(RewritePattern): + """Replace each pre-discovered subgraph with a fused op, walker-driven. + + The pattern matches on each subgraph's *anchor* — its topologically-first op + (``sg.ops[0]``). Because discovered subgraphs are op-disjoint, anchors are + unique and survive until their own subgraph is fused. A forward walk visits + partition anchors in dependency order, so for subgraphs split across the + 64-input limit (where a later partition consumes an earlier one's output), + ``refresh_inputs()`` picks up the already-redirected operands — matching the + behavior of the previous ahead-of-time discover-then-replace loop. + """ + + def __init__( + self, + subgraphs: List[FusibleSubgraph], + metadata: Dict, + log_warning: Optional[Callable[[str], None]] = None, + ): + super().__init__() + # Anchor op -> subgraph. Discovered subgraphs are op-disjoint, so each + # anchor maps to exactly one subgraph. + self._by_anchor: Dict[Operation, FusibleSubgraph] = { + sg.ops[0]: sg for sg in subgraphs if sg.ops + } + self._metadata = metadata + self._log_warning = log_warning + self.num_replaced = 0 + self.num_skipped_low_rank = 0 + + def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter) -> None: + sg = self._by_anchor.get(op) + if sg is None: + return + + # Skip subgraphs where all inputs are 1D or lower — these are pure + # weight-space ops (e.g., weight + 1.0) that don't benefit from fusion + # and the row-based Triton kernel can't handle them correctly. + if _max_input_rank(sg) < 2 or _min_output_rank(sg) < 2: + self.num_skipped_low_rank += 1 + return + + try: + # Refresh inputs: earlier subgraph replacements may have redirected + # operands via SSAValue.replace_by(), making the inputs list computed + # at discovery time stale. + sg.refresh_inputs() + kernel_fn = generate_kernel_from_subgraph(sg) + if kernel_fn is None: + return + sg_hash = KernelCache.hash_subgraph(sg) + replace_subgraph_with_fused_op( + sg, kernel_fn, sg_hash, self._metadata, rewriter=rewriter + ) + self.num_replaced += 1 + except (ValueError, NotImplementedError) as e: + if self._log_warning is not None: + self._log_warning(f"Skipping subgraph (unsupported pattern): {e}") + + +def run_fusion( + mlir_module: ModuleOp, + metadata: Dict, + log_warning: Optional[Callable[[str], None]] = None, +) -> FusionStats: + """Discover fusible subgraphs and replace each with a generated fused op. + + Args: + mlir_module: The module to fuse in place. + metadata: The ``FXToMLIRConverter`` metadata side-table; one entry is + added per fused op so MLIR-to-FX can reconstruct the kernel call. + log_warning: Optional callback for per-subgraph skip diagnostics. + + Returns: + A :class:`FusionStats` with discovery/replacement counts. + """ + subgraphs = discover_fusible_subgraphs(mlir_module) + if not subgraphs: + return FusionStats() + + pattern = _FuseSubgraphPattern(subgraphs, metadata, log_warning) + # apply_recursively=False: a single forward pass; each anchor is matched once + # and the inserted fused op is never re-matched. + walker = PatternRewriteWalker( + GreedyRewritePatternApplier([pattern]), + apply_recursively=False, + ) + walker.rewrite_module(mlir_module) + + return FusionStats( + num_subgraphs=len(subgraphs), + num_replaced=pattern.num_replaced, + num_skipped_low_rank=pattern.num_skipped_low_rank, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.py b/tensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.py index 73de63e3bf00..c48387ef7d25 100644 --- a/tensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.py +++ b/tensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.py @@ -21,9 +21,10 @@ Triton kernel. """ -from typing import Any, Callable, Dict +from typing import Any, Callable, Dict, Optional from xdsl.dialects.builtin import StringAttr +from xdsl.pattern_rewriter import InsertPoint, PatternRewriter from ..dialect import AdOpaque from .subgraph_discovery import FusibleSubgraph @@ -34,6 +35,7 @@ def replace_subgraph_with_fused_op( kernel_fn: Callable, sg_hash: str, metadata: Dict[str, Dict[str, Any]], + rewriter: Optional[PatternRewriter] = None, ) -> None: """Replace a subgraph's ops in the MLIR block with a single fused AdOpaque. @@ -44,6 +46,12 @@ def replace_subgraph_with_fused_op( sg_hash: The subgraph hash used in the torch.ops registration name. metadata: The FXToMLIRConverter metadata side-table. A new entry is added for the fused op so MLIR-to-FX can reconstruct it. + rewriter: Optional xDSL ``PatternRewriter``. When provided (the + walker-driven :func:`run_fusion` path), op insertion and erasure go + through the rewriter so SSA rewiring is tracked and validated by + xDSL's pass infrastructure. When ``None`` (standalone callers, + e.g. unit tests), the same mutations are applied directly on the + block with the default safe-erase use-def check. """ import torch from xdsl.dialects.builtin import TensorType @@ -111,17 +119,29 @@ def replace_subgraph_with_fused_op( # An external input is produced AFTER the first subgraph op. # Insert the fused op just after that input producer. anchor_op = list(block.ops)[latest_input_pos] - block.insert_op_after(fused_op, anchor_op) + if rewriter is not None: + rewriter.insert_op(fused_op, InsertPoint.after(anchor_op)) + else: + block.insert_op_after(fused_op, anchor_op) else: # Normal case: all inputs are available before the subgraph. - block.insert_op_before(fused_op, subgraph.ops[0]) + if rewriter is not None: + rewriter.insert_op(fused_op, InsertPoint.before(subgraph.ops[0])) + else: + block.insert_op_before(fused_op, subgraph.ops[0]) # Replace each subgraph output's uses with the corresponding fused op output for i, out_val in enumerate(subgraph.outputs): out_val.replace_by(fused_op.outputs[i]) - # Erase the original subgraph ops in reverse order. - # safe_erase=False skips the use-check, which is needed because internal - # operand references between subgraph ops may still exist at erase time. + # Erase the original subgraph ops in reverse topological order. + # ``subgraph.ops`` is topologically sorted (producers before consumers), so + # iterating in reverse erases consumers first. Combined with the + # ``replace_by`` above (which redirects every external use of a subgraph + # output), each op has no remaining uses by the time it is erased, so the + # default ``safe_erase=True`` use-def check passes. for op in reversed(subgraph.ops): - block.erase_op(op, safe_erase=False) + if rewriter is not None: + rewriter.erase_op(op) + else: + block.erase_op(op) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py b/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py index cd9ce643f1df..1595b23f537c 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/mlir_elementwise_fusion.py @@ -82,11 +82,8 @@ def _apply( self._log_warning("xDSL not installed, skipping MLIR elementwise fusion") return gm, TransformInfo(skipped=True) - from ...mlir.codegen.kernel_cache import KernelCache - from ...mlir.codegen.triton_emitter import generate_kernel_from_subgraph from ...mlir.decompose import run_decomposition - from ...mlir.fusion.subgraph_discovery import discover_fusible_subgraphs - from ...mlir.fusion.subgraph_replace import replace_subgraph_with_fused_op + from ...mlir.fusion.fuse import run_fusion from ...mlir.fx_to_mlir import FXToMLIRConverter from ...mlir.mlir_to_fx import MLIRToFXConverter @@ -104,60 +101,16 @@ def _apply( num_decomposed = run_decomposition(mlir_module) self._log_info(f"Decomposed {num_decomposed} high-level ops into primitives") - # Step 3: Discover fusible subgraphs - subgraphs = discover_fusible_subgraphs(mlir_module) + # Step 3+4: Discover fusible subgraphs and replace each with a generated + # fused op. Both decomposition and fusion now go through xDSL's + # PatternRewriteWalker (see run_decomposition / run_fusion). + stats = run_fusion(mlir_module, converter.metadata, log_warning=self._log_warning) self._log_info( - f"Discovered {len(subgraphs)} fusible subgraphs " - f"(total ops: {sum(len(sg.ops) for sg in subgraphs)})" + f"Discovered {stats.num_subgraphs} fusible subgraphs; " + f"replaced {stats.num_replaced} (skipped {stats.num_skipped_low_rank} low-rank)" ) - if not subgraphs: - return gm, TransformInfo( - skipped=False, num_matches=0, is_clean=True, has_valid_shapes=True - ) - - # Step 4: Generate Triton kernels and replace subgraphs in MLIR. - # Skip subgraphs where all inputs are 1D or lower — these are pure - # weight-space ops (e.g., weight + 1.0) that don't benefit from fusion - # and the row-based Triton kernel can't handle them correctly. - from xdsl.dialects.builtin import TensorType as _TT - - def _max_input_rank(sg): - return max( - (len(inp.type.get_shape()) for inp in sg.inputs if isinstance(inp.type, _TT)), - default=0, - ) - - def _min_output_rank(sg): - return min( - (len(out.type.get_shape()) for out in sg.outputs if isinstance(out.type, _TT)), - default=0, - ) - - num_replaced = 0 - num_skipped = 0 - for sg in subgraphs: - if _max_input_rank(sg) < 2 or _min_output_rank(sg) < 2: - num_skipped += 1 - continue - try: - # Refresh inputs: earlier subgraph replacements may have - # redirected operands via SSAValue.replace_by(), making the - # inputs list computed at discovery time stale. - sg.refresh_inputs() - kernel_fn = generate_kernel_from_subgraph(sg) - if kernel_fn is not None: - sg_hash = KernelCache.hash_subgraph(sg) - replace_subgraph_with_fused_op(sg, kernel_fn, sg_hash, converter.metadata) - num_replaced += 1 - except (ValueError, NotImplementedError) as e: - self._log_warning(f"Skipping subgraph (unsupported pattern): {e}") - - self._log_info( - f"Replaced {num_replaced}/{len(subgraphs)} subgraphs (skipped {num_skipped} low-rank)" - ) - - if num_replaced == 0: + if stats.num_replaced == 0: return gm, TransformInfo( skipped=False, num_matches=0, is_clean=True, has_valid_shapes=True ) @@ -174,7 +127,7 @@ def _min_output_rank(sg): # happens later at cache_init). return new_gm, TransformInfo( skipped=False, - num_matches=num_replaced, + num_matches=stats.num_replaced, is_clean=False, has_valid_shapes=False, ) diff --git a/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py b/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py index 398950853709..9daa32ce7ebf 100644 --- a/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py +++ b/tests/unittest/auto_deploy/singlegpu/mlir/test_elementwise_fusion_e2e.py @@ -497,6 +497,134 @@ def test_moe_ep_mask_fusion_discovery(): assert len(subgraphs[0].ops) == 3, f"Expected 3 ops, got {len(subgraphs[0].ops)}" +# --------------------------------------------------------------------------- +# Walker-driven run_fusion (PatternRewriteWalker) — parity with the manual +# discover-then-replace loop. +# --------------------------------------------------------------------------- + + +def test_run_fusion_replaces_subgraph_and_records_metadata(): + """run_fusion discovers + replaces a subgraph via the walker, no CUDA needed. + + Exercises the xDSL PatternRewriteWalker path: the AdAdd+AdMul subgraph is + fused into a single AdOpaque, the original primitives are erased, and a + fused-op metadata entry is recorded with the tuple "val" contract. + """ + from tensorrt_llm._torch.auto_deploy.mlir.dialect import AdOpaque + from tensorrt_llm._torch.auto_deploy.mlir.fusion.fuse import run_fusion + + hidden = 128 + mlir_mod = _build_single_output_add_mul_module(hidden) + + metadata = {} + stats = run_fusion(mlir_mod, metadata) + + assert stats.num_subgraphs == 1 + assert stats.num_replaced == 1 + assert stats.num_skipped_low_rank == 0 + + block = mlir_mod.body.block + # The fused AdOpaque replaced the AdAdd/AdMul primitives. + assert sum(isinstance(op, AdOpaque) for op in block.ops) == 1 + assert not any(isinstance(op, (AdAdd, AdMul)) for op in block.ops) + + # Exactly one fused-op metadata entry, with the tuple "val" contract. + assert len(metadata) == 1 + (val_meta,) = (m["val"] for m in metadata.values()) + assert isinstance(val_meta, tuple) and len(val_meta) == 1 + assert tuple(val_meta[0].shape) == (2, hidden) + + +def test_run_fusion_skips_low_rank_subgraph(): + """run_fusion reports low-rank subgraphs as skipped, not replaced.""" + from xdsl.dialects.builtin import Region as _Region + + from tensorrt_llm._torch.auto_deploy.mlir.fusion.fuse import run_fusion + + # Two chained ops on 1D (weight-space) tensors → discovered but low-rank. + tw = TensorType(BFloat16Type(), [128]) + block = Block() + a = block.insert_arg(tw, 0) + b = block.insert_arg(tw, 1) + add_op = AdAdd.build(operands=[a, b], result_types=[tw]) + block.add_op(add_op) + mul_op = AdMul.build(operands=[add_op.output, b], result_types=[tw]) + block.add_op(mul_op) + block.add_op(AdGraphOutput.build(operands=[[mul_op.output]])) + mlir_mod = ModuleOp(_Region([block])) + + metadata = {} + stats = run_fusion(mlir_mod, metadata) + + assert stats.num_subgraphs == 1 + assert stats.num_replaced == 0 + assert stats.num_skipped_low_rank == 1 + assert metadata == {} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_run_fusion_full_fx_roundtrip(): + """FX → MLIR → decompose → run_fusion (walker) → MLIR → FX matches eager. + + Walker-path counterpart of ``test_full_fx_roundtrip_with_replacement``, + which drives the same model through the manual discover-then-replace loop. + """ + import tensorrt_llm._torch.auto_deploy.custom_ops.normalization.rms_norm # noqa: F401 + import tensorrt_llm._torch.auto_deploy.mlir # noqa: F401 + from tensorrt_llm._torch.auto_deploy.mlir.fusion.fuse import run_fusion + from tensorrt_llm._torch.auto_deploy.mlir.fx_to_mlir import FXToMLIRConverter + from tensorrt_llm._torch.auto_deploy.mlir.mlir_to_fx import MLIRToFXConverter + + hidden = 128 + + class AddNormModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + ) + self.eps = 1e-5 + + def forward(self, x, residual): + added = x + residual + norm = torch.ops.auto_deploy.torch_rmsnorm(added, self.weight, self.eps) + return norm, added + + model = AddNormModel() + x = torch.randn(2, 8, hidden, device="cuda", dtype=torch.bfloat16) + res = torch.randn_like(x) + + from torch.export import Dim + + from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm + + gm = torch_export_to_gm( + model, + args=(x, res), + dynamic_shapes=({0: Dim.DYNAMIC}, {0: Dim.DYNAMIC}), + clone=True, + ) + + converter = FXToMLIRConverter(gm) + mlir_module = converter.convert() + + num_decomposed = run_decomposition(mlir_module) + assert num_decomposed >= 1 + + stats = run_fusion(mlir_module, converter.metadata) + assert stats.num_replaced >= 1 + + back_converter = MLIRToFXConverter(gm) + new_gm = back_converter.convert(mlir_module, converter.metadata) + + result = new_gm(x.clone(), res.clone()) + ref = model(x.clone(), res.clone()) + + assert len(result) == len(ref) + torch.testing.assert_close(result[0], ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(result[1], ref[1], atol=1e-2, rtol=1e-2) + + def test_moe_ep_mask_kernel_generation(): """Triton kernel can be generated for the mixed-type floordiv+eq+mul pattern.""" mlir_mod = _build_moe_ep_mask_module()