[Feature] add lower-trace support for debugging#2470
Open
erhsh wants to merge 1 commit into
Open
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Contributor
📝 WalkthroughWalkthroughAdds a new ChangesIR Lower Trace Feature
Sequence DiagramsequenceDiagram
participant User
participant tilelang.__init__
participant lower_trace.enable
participant tvm.Pass.__call__
participant _wrap_codegen_ffi
participant generate_html
User->>tilelang.__init__: import tilelang (TL_LOWER_TRACE set)
tilelang.__init__->>lower_trace.enable: enable(mode, trace_dir)
lower_trace.enable->>tvm.Pass.__call__: monkey-patch → _traced_pass_call
lower_trace.enable->>target.build.*: monkey-patch → _wrap_codegen_ffi
User->>tvm.Pass.__call__: pass(mod)
tvm.Pass.__call__->>_traced_pass_call: snapshot before IR
_traced_pass_call->>tvm.Pass.__call__: call original pass
_traced_pass_call->>_traced_pass_call: snapshot after IR, diff, save LowerRecord
User->>target.build.*: build(mod, target)
target.build.*->>_wrap_codegen_ffi: capture TIR, invoke original build
_wrap_codegen_ffi->>_wrap_codegen_ffi: three-way baseline/working/latest comparison
_wrap_codegen_ffi->>_traced_pass_call: append STATUS_CODEGEN record
User->>generate_html: atexit / finally block
generate_html->>generate_html: render all LowerRecords → report.html
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. 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: 6
🤖 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 `@testing/python/debug/test_lower_trace.py`:
- Around line 9-20: The top-level TL_LOWER_TRACE cleanup in the test module
removes the caller’s environment settings for the rest of the pytest process;
update the import guard around the initial tilelang import so it temporarily
clears TL_LOWER_TRACE and TL_LOWER_TRACE_DIR, then restores the previous values
immediately after import. Use the existing import block and the _isolate_env
fixture in test_lower_trace.py as the anchor points, and keep the rest of the
test behavior unchanged.
In `@tilelang/tools/lower_trace/__init__.py`:
- Around line 108-123: The terminal-mode trace handling in the pass loop around
p(mod) re-raises too early, so the failing step and captured pre-pass IR are
never printed when an exception occurs. Update the exception path in the loop
that builds results and calls p(mod) to emit the terminal header or an error
stanza from the except block before re-raising, ensuring mode="terminal" still
shows which pass failed and the available IR context.
- Around line 168-190: The HTML report generation in the `finally` block can
overwrite the original lowering exception if `generate_html` fails, so preserve
the first error from the lowering flow and treat report writing as best-effort.
Update the `finally` logic around `results`, `generate_html`, and the
`LowerRecord` creation so any exception from HTML writing is caught and logged
without replacing the original failure, ensuring the real pass error remains the
one that propagates.
In `@tilelang/tools/lower_trace/core.py`:
- Around line 313-320: Honor terminal-only mode by preventing implicit
filesystem writes when tracing is set to terminal. Update
`_get_codegen_output_path()` in `core.py` so it only falls back to
`<script_dir>/codegen.cpp` for HTML/both persistence modes, while still
returning any explicit override; then adjust `_save_raw_files()` and the related
run-directory/file creation path to skip creating output directories/files
unless persistence is enabled. Keep explicit `codegen_output` support working
even in terminal mode, and use the existing trace-mode helpers around
`_is_trace_enabled()` / `_ensure_script_dir()` to separate terminal-only from
file-writing behavior.
In `@tilelang/tools/lower_trace/html.py`:
- Around line 1014-1026: The expand/collapse handlers in expandAll and
collapseCtx assume a table always exists inside the current pass-section, but
_make_diff_html() can render only a newline-only p while rec.changed is true, so
the toolbar may invoke these functions with no table present. Add a null check
before querying or using table (and skip the row updates when absent), or
prevent rendering the expand/collapse toolbar when no table is produced, so the
handlers fail safely in the affected diff rendering paths.
- Line 1365: The DOM ID generation in html.py is using rec.phase directly in
sec-* identifiers, which can break markup and selectors if the phase contains
quotes or CSS-special characters. Update the ID generation and any related
selector usage in the affected rendering helpers (including the record/detail
sections and the later lookup code) to derive a stable sanitized DOM-safe ID
from rec.phase, while keeping the original phase only for display text. Make
sure the same safe-ID helper is used consistently wherever these IDs are created
or referenced.
🪄 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: Pro
Run ID: 46ec0d80-f612-4e1a-bf6e-5edeecd33645
⛔ Files ignored due to path filters (1)
docs/_static/img/lower_trace_html.pngis excluded by!**/*.png
📒 Files selected for processing (8)
docs/tutorials/debug_tools_for_tilelang.mdtesting/python/debug/test_lower_trace.pytilelang/__init__.pytilelang/tools/__init__.pytilelang/tools/lower_trace/__init__.pytilelang/tools/lower_trace/core.pytilelang/tools/lower_trace/diff.pytilelang/tools/lower_trace/html.py
Contributor
Author
|
@LeiWang1999 please help review & merge, thx |
erhsh
changed the title
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
1 participant
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.
IR Lower Trace — Feature Overview
1. One-line Summary
A zero-intrusion compilation debugging tool: once enabled, it automatically captures the IR before and after every pass (and the final codegen) in the compilation pipeline, renders a per-pass diff in the terminal and as a self-contained HTML report, and supports an "edit the generated codegen source → recompile with edits injected" workflow for verification.
2. Why It's Needed
TileLang programs are lowered through a sequence of compiler passes, each of which may rewrite the IR. When the final result is unexpected, it is hard to tell which pass changed the IR incorrectly, which backend stage a pass belongs to, or what the final generated C/CUDA/HIP source looks like — let alone quickly patch the generated kernel for experimentation without touching the TileLang source. This tool closes those observability gaps: off → zero overhead; on → the whole pipeline is visible.
3. Implementation
Located under tilelang/tools/lower_trace/, the tool consists of three parts: hook installation & record collection (core), unified diff computation (diff), and self-contained HTML rendering (html).
Pass.__call__to capture before/after IR for each pass; interceptPassPipeline.lower(new arch) / phase functions (legacy arch, auto-discovered via AST scan + bytecode inspection) to tag each pass with its pipeline stage; intercept codegen FFIs to capture the final TIR → source step.run2_/run3_prefixes and merged into a single report for side-by-side comparison.PATCHED/REGENERATED/SYNCED/CONFLICT, so user edits are never silently overwritten (requires a source-compiling backend:nvrtc/cython/cutedsl).inspect_source/get_source; path components are sanitized against directory escape; disk writes are best-effort (failures warn, never abort compilation).4. How to Use
Environment variable (simplest):
TL_LOWER_TRACE_DIRsets the artifact root (default./tmp/lower_trace_dir); the report is at<dir>/<script>/report.html.Programmatic API:
enable()is idempotent (optionalmode/trace_dir/codegen_output);reset()clears records but keeps the hook (per-kernel reports in one process);disable()restores all hooks.edit-recompile flow: run once with tracing on → edit the on-disk
codegen.cpp→ rerun to compile the edited source; if a TileLang source change regenerates codegen that conflicts with your edits, the working copy is auto-backed-up to.bakfor recovery.HTML report features: sidebar with per-pass navigation + status dots (changed / no-op / failed / codegen) + line stats, stage filtering, side-by-side diff (word-level highlighting, collapsible context),
j/kkeyboard navigation, dark/light theme, and failed-pass error boxes showing the exception alongside the pre-crash IR.5. Impact on the Existing Framework (non-intrusive, easy to merge)
This tool is added in a purely additive, non-intrusive way — no existing framework behavior is changed, and merge risk is low:
TL_LOWER_TRACEis unset, the loading logic in tilelang/init.py is skipped entirely — no import, no patching, no runtime cost; tilelang/tools/init.py lazily loads the submodule via__getattr__, leaving the normal startup path untouched.disable()restoresPass.__call__,PassPipeline.lower, and codegen FFIs to their originals and clears all state — no trace left.tilelang/tools/lower_trace/module plus tests and docs; the only edits to existing source are an env-gated loading block and a lazy-load entry in__init__.py— no compilation-core logic is touched.In short, the feature provides full-pipeline observability without touching the framework's compilation core, is off by default with zero production impact, and is suitable for direct merge.
Appendix: Artifact directory layout
Summary
lower_tracedebugging workflow for TileLang lowering, with hook-based IR capture before/after passes, per-pass diffs, terminal output, and self-contained HTML reporting.TL_LOWER_TRACEand lazy module loading throughtilelang.tools.Testing
testing/python/debug/test_lower_trace.pywith broad regression coverage for the new tracing flow.