[Debug][TIR] Drive Pass Visualizer with PassInstrument

[LeiWang1999]

Summary

• Replace the Pass Visualizer duplicate CUDA prologue list with a TVM PassInstrument attached to the canonical prologue.
• Keep structure-tree reports synchronized with target-dependent and pass-config-dependent compilation behavior.

Changes

• Capture before/after structure trees for each top-level pass while folding nested implementation passes into their parent stage.
• Run semantic checks and the canonical CUDA prologue with normalized PrimFunc and JIT pass configurations from the analyzed kernel.
• Add stable PassInfo names for Python-defined passes visible in the instrumented report.
• Remove the fixed build_pass_stages helper and update the public exports, documentation, and regression coverage.
• Cover real pipeline observation, conditional pass configuration, nested callbacks, failure cleanup, and HTML/text emission.

Validation

• ./format.sh
• python -m pytest testing/python/debug/test_pass_visualizer.py testing/python/transform/test_tilelang_transform_decouple_type_cast.py -x -q (26 passed, 1 skipped)
• python -m pytest testing/python/debug/test_lower_trace.py testing/python/debug/test_pass_diff.py -x -q (76 passed)

Notes

• The viewer remains focused on CUDAPassPipelineBodyPrologue; full-pipeline tracing remains available through TL_LOWER_TRACE.

Summary

• Refactored the Pass Visualizer to use StructureTreePassInstrument on the canonical CUDAPassPipelineBodyPrologue.
• Captured before/after structure trees for top-level passes.
• Grouped nested implementation passes under their parent stages.
• Applied normalized JIT and PrimFunc pass configurations from the analyzed kernel.
• Added stable names for Python-defined passes.
• Removed build_pass_stages.
• Added capture_structure, PassStructureRecord, and StructureTreePassInstrument exports.
• Added explicit names for several TileLang passes.
• Updated documentation and CLI descriptions to identify CUDA support and canonical pipeline tracing.
• Added regression coverage for conditional ordering, nested callbacks, change detection, failure cleanup, dynamic passes, and pass configurations.

Validation

• 26 tests passed, 1 skipped.
• 76 tests passed.
• Formatting and targeted test suites completed successfully.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: c7bf2a0e-defa-4f9d-a9cc-76abc03a840c

📥 Commits

Reviewing files that changed from the base of the PR and between 42dabfe and 59c38ed.

📒 Files selected for processing (5)

• docs/tutorials/debug_tools_for_tilelang.md
• testing/python/debug/test_pass_visualizer.py
• tilelang/tools/pass_visualizer/README.md
• tilelang/tools/pass_visualizer/core.py
• tilelang/tools/pass_visualizer/viewer.py

🚧 Files skipped from review as they are similar to previous changes (3)

• tilelang/tools/pass_visualizer/README.md
• tilelang/tools/pass_visualizer/viewer.py
• docs/tutorials/debug_tools_for_tilelang.md

---

📝 Walkthrough

Walkthrough

The pass visualizer now observes the canonical CUDA lowering prologue through TVM PassInstrument. It records ordered top-level passes, structure snapshots, nested-pass behavior, failures, dynamic pipeline changes, and JIT configurations.

Changes

Pass visualizer instrumentation

Layer / File(s)	Summary
Instrumented pass capture tilelang/tools/pass_visualizer/core.py, tilelang/tools/pass_visualizer/__init__.py, tilelang/analysis/layout_visual.py, tilelang/transform/*	Replaces manual CUDA stage construction with StructureTreePassInstrument, ordered PassStructureRecord data, structure capture, callback diagnostics, and explicit pass names.
Viewer pipeline integration tilelang/tools/pass_visualizer/viewer.py	Validates CUDA targets, applies normalized pass configurations, runs PreLowerSemanticCheck, and renders instrument-captured pass records and diffs from CUDAPassPipelineBodyPrologue.
Validation and documentation testing/python/debug/test_pass_visualizer.py, docs/tutorials/debug_tools_for_tilelang.md, tilelang/tools/pass_visualizer/README.md	Tests canonical pass recording, nested-pass filtering, cleanup, dynamic passes, target rejection, and JIT configuration. Documentation describes the CUDA-specific instrumented flow.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant build_pass_data
  participant CUDAPassPipelineBodyPrologue
  participant StructureTreePassInstrument
  participant HTMLReport
  CLI->>build_pass_data: request CUDA pass visualization
  build_pass_data->>CUDAPassPipelineBodyPrologue: run configured lowering prologue
  CUDAPassPipelineBodyPrologue->>StructureTreePassInstrument: emit pass callbacks
  StructureTreePassInstrument->>build_pass_data: return ordered pass records and snapshots
  build_pass_data->>HTMLReport: render pass names and structure diffs

Loading

Possibly related PRs

• tile-ai/tilelang#2375: Earlier pass instrumentation for capturing before/after IR and visualizing pass changes.
• tile-ai/tilelang#2449: Previous manual build_pass_stages implementation replaced by this instrumented flow.
• tile-ai/tilelang#2859: Earlier changes that mirrored the CUDA prologue in build_pass_stages.

Suggested reviewers: siriusneo, yongqi-zhuo

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: refactoring the TIR Pass Visualizer to use PassInstrument.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Copilot]

Pull request overview

This PR refactors the TIR Pass Visualizer to observe the real CUDA lowering prologue via a TVM PassInstrument, eliminating the need to maintain a duplicate prologue pass list and keeping pass ordering/conditionality aligned with target and pass-config behavior.

Changes:

• Introduces StructureTreePassInstrument to capture before/after structure-tree snapshots for top-level passes (folding nested implementation passes).
• Updates the visualizer CLI to run the canonical CUDAPassPipelineBodyPrologue under instrumented PassContext, including normalized pass configs and semantic checks.
• Stabilizes Python-defined pass naming via explicit prim_func_pass(..., name=...), and updates docs/exports/tests accordingly.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tilelang/transform/decouple_type_cast.py	Assigns a stable PassInfo name for the PrimFunc pass to improve instrumented reporting.
tilelang/transform/add_bufstore_wrapper.py	Assigns a stable PassInfo name for the PrimFunc pass to improve instrumented reporting.
tilelang/analysis/layout_visual.py	Assigns a stable PassInfo name for the analysis PrimFunc pass.
tilelang/tools/pass_visualizer/core.py	Adds structure-tree capture helper + StructureTreePassInstrument, removes duplicate prologue pass list builder.
tilelang/tools/pass_visualizer/viewer.py	Switches from a locally-built pass list to instrumenting the canonical CUDA prologue and emitting diffs from instrument snapshots.
tilelang/tools/pass_visualizer/init.py	Updates public exports to reflect the new instrument-based API and removes build_pass_stages.
tilelang/tools/pass_visualizer/README.md	Updates documentation to describe instrumented real-prologue observation and nested-pass folding.
testing/python/debug/test_pass_visualizer.py	Reworks tests to validate instrument behavior, config precedence, injected passes, and end-to-end HTML/text emission.
docs/tutorials/debug_tools_for_tilelang.md	Updates tutorial text to match the new PassInstrument-based behavior.

Suppressed comments (1)

tilelang/tools/pass_visualizer/core.py:453

• PassStructureRecord.after_ir is currently set to str(mod), which can be very large and can cause the changed flag to disagree with the structure-tree diff shown in the UI. Prefer reusing the captured structure lines for after_ir so changed reflects the structure-tree snapshot and avoids holding full TVMScript strings.

        self.records.append(
            PassStructureRecord(
                name=name,
                sequence=frame.sequence,
                before_ir=frame.before_ir,
                after_ir=str(mod),
                before_lines=frame.before_lines,
                after_lines=capture_structure(mod),
            )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SiriusNEO]

Generally I think it's good to replace the original impl with the PassInstrument approach. Left some comments