[None][refactor] AutoDeploy: make MLIR elementwise fusion pipeline xDSL-native

.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=<HF_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 <changed_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.

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.

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,
+    )