IR Lower Trace

IR Lower Trace shows how TileLang’s TIR changes as compiler passes run. It transparently captures the IR before and after every pass in the compilation pipeline — including the final codegen step that produces C/CUDA/HIP source — and renders a human-readable diff report in the terminal and/or a self-contained HTML page. No code changes are required: enabling a single environment variable is enough.

IR Lower Trace is the recommended successor to Pass Diff. Compared with the older TILELANG_PASS_DIFF tool, it adds:

• Phase context — each pass is tagged with its pipeline phase (e.g. pipeline_c, phase1_...), so you can tell which backend stage a pass belongs to.

• Codegen capture — the final TIR-to-source lowering is recorded, and the generated C/CUDA/HIP code is dropped to disk for inspection or editing.

• Edit-and-recompile workflow — edit the generated codegen source on disk and rerun; your edits are injected back into compilation (with conflict detection).

• Multi-run accumulation — repeated compilations in the same process are tagged with run2_, run3_, … prefixes, so you can diff across runs.

• Raw .tir dumps — before/after IR for every pass is written to disk, keyed by phase and pass index.

• Crash-safe incremental HTML — the report is flushed after every pass, so partial results survive even if the process crashes.

• Enhanced HTML report — sidebar pass navigation, status dots, phase tabs, j/k keyboard navigation, Shift+E global expand, F7 manual alignment, dark/light theme.

Observe the Lowering Pipeline

The simplest way to enable IR Lower Trace is to set the TL_LOWER_TRACE environment variable before running your script:

# HTML report (default when set to 1/on/true/yes)
TL_LOWER_TRACE=1 python3 my_script.py

# Colored diff printed to the terminal only
TL_LOWER_TRACE=terminal python3 my_script.py

# Both terminal output and HTML report
TL_LOWER_TRACE=both python3 my_script.py

# Disabled (default — zero overhead, no patching)
python3 my_script.py

TileLang reads this setting when tilelang is imported and installs a process- wide hook around TVM pass execution. Setting the variable after importing TileLang does not enable the hook for that process.

When HTML output is enabled, a stable symlink <script_dir>/report.html is maintained and points to the latest run’s report. Open it directly in a browser:

# typical location
open tmp/lower_trace_dir/my_script/report.html

Variable	Description	Default
TL_LOWER_TRACE	Enable tracing. Values: 0/off/false/no (off), 1/on/true/yes (→ html), terminal, html, both	off
TL_LOWER_TRACE_DIR	Base output directory for all trace artifacts	./tmp/lower_trace_dir

Output Directory Structure

A single run produces the following layout under TL_LOWER_TRACE_DIR:

<TL_LOWER_TRACE_DIR>/
└── <script_name>/                      # derived from sys.argv[0], e.g. "my_script"
    ├── report.html                # symlink → latest run's report
    ├── codegen.cpp                     # generated codegen source (editable, see below)
    ├── codegen.cpp.original            # baseline snapshot for edit/recompile workflow
    ├── codegen.cpp.latest              # actual codegen output of the most recent run
    └── .run_records/
        └── run_<YYYYMMDD_HHMMSS_ffffff>_<pid>/
            ├── report.html        # this run's full report
            ├── pipeline_c/             # one subdir per phase (example)
            │   ├── 00_BindTarget_before.tir
            │   ├── 00_BindTarget_after.tir
            │   ├── 01_Simplify_before.tir
            │   └── 01_Simplify_after.tir
            ├── phase2_optimize/        # another phase (illustrative)
            │   └── ...
            ├── codegen/
            │   ├── 42_codegen_before.tir
            │   └── 42_codegen_after.cpp
            └── unscoped/               # passes outside any pipeline window
                └── ...

Each phase gets its own subdirectory; passes are numbered globally (00_, 01_, …) so ordering is unambiguous. Codegen records write the after artifact as .cpp (the generated source), while ordinary passes use .tir for both sides.

Programmatic API

For fine-grained control, IR Lower Trace exposes two layers of API.

One-shot: lower_trace()

Diff a fixed chain of passes against an IR module without installing any global hook:

from tilelang.tools import lower_trace as lt
from tilelang import tvm
import tilelang.transform as transform

# Diff a single pass
results = lt.lower_trace(func, transform.Simplify(), mode="terminal")

# Diff a named chain, write an HTML report
results = lt.lower_trace(
    func,
    [
        ("Annotate",   tvm.tirx.transform.AnnotateDeviceRegions()),
        ("Split",      tvm.tirx.transform.SplitHostDevice()),
        ("ThreadSync", transform.ThreadSync("shared")),
    ],
    mode="both",
    html_path="my_diff.html",
)

Parameter	Description	Default
func_or_mod	A PrimFunc or IRModule to run passes on	Required
passes	A single pass, a list of passes, or a list of (name, pass) tuples	Required
mode	"terminal", "html", or "both"	terminal
context	Number of context lines in the unified diff	3
html_path	Output path for the HTML report	lower_trace_report.html

Returns a list[dict] with one entry per pass step, each containing name, before_script, after_script, diff_lines, insertions, deletions, and changed.

Global hook: enable() / disable() / reset()

To trace the entire compilation pipeline of a real kernel (what the environment variable does, but programmatically):

from tilelang.tools import lower_trace as lt

# Enable tracing for the rest of the process.
lt.enable(mode="both")

# ... run tilelang.compile() / kernel compilation ...

All three parameters of lt.enable() are optional — mode, trace_dir, and codegen_output fall back to the TL_LOWER_TRACE / TL_LOWER_TRACE_DIR env vars (or sensible defaults) when omitted.

Parameter	Description	Default
mode	Force a trace mode: "terminal", "html", "both", or None to disable.	TL_LOWER_TRACE env var
trace_dir	Base output directory.	TL_LOWER_TRACE_DIR, then ./tmp/lower_trace_dir
codegen_output	Path to save the generated codegen source (enables the edit-recompile workflow). Pass None explicitly to suppress. Defaults to <script_dir>/codegen.cpp in html/both mode; None (no file) in terminal mode unless overridden.

enable() is idempotent — calling it multiple times is safe.

When to use reset() and disable()

Both are optional and only needed in specific scenarios:

Function	When to call	What it does
reset()	Compiling multiple kernels in the same process and you want each kernel’s report to start fresh (instead of accumulating into one combined report)	Clears collected records while keeping the hook active. Without it, records accumulate across compilations, tagged with run2_, run3_, … prefixes — which is desirable if you want to compare runs side by side.
disable()	You want to disable tracing for subsequent compilations within the same process (e.g. a long-running service that only traces the first kernel)	Restores the original Pass.__call__, PassPipeline.lower, and codegen FFIs, and clears all state.

from tilelang.tools import lower_trace as lt

lt.enable(mode="both")

# First kernel — traced.
kernel1 = tilelang.compile(func_a)

# Optional: clear records so kernel2 gets its own clean report.
# Omit this line if you prefer a combined multi-run report.
lt.reset()

# Second kernel — traced (into a fresh report if lt.reset() was called).
kernel2 = tilelang.compile(func_b)

# Optional: disable tracing for any further compilations.
lt.disable()

Note

If neither lt.reset() nor lt.disable() is called, tracing stays active for the lifetime of the process and the final HTML report is generated automatically at exit. This is the simplest workflow and is sufficient for most one-off scripts.

HTML Report Features

The HTML report is a single self-contained file (no external assets) providing:

• Sidebar with per-pass navigation, status dots (● changed / ○ no-op / ✕ failed / ◆ codegen), and +/− line-count statistics. Collapsible and drag-resizable.

• Phase tabs to filter passes by pipeline phase.

• Summary bar with clickable filter badges (changed / failed / codegen).

• Side-by-side diff with GitHub-style coloring, word-level inline highlighting, and collapsible context (↑↓ Expand buttons reveal hidden equal lines).

• Keyboard navigation — j/k to move between passes, Shift+E to expand/collapse all, F7 for Beyond-Compare-style manual alignment.

• Dark/Light theme toggle persisted via localStorage.

• Copy buttons for before/after IR of any pass.

• Error boxes — a failed pass shows its exception message alongside the IR before the crash.

IR Lower Trace HTML report.

Codegen Source Capture & Edit-Recompile Workflow

When tracing is enabled, the final codegen step (TIR → C/CUDA/HIP/…) is intercepted. The generated source is written to <script_dir>/codegen.cpp (or the path passed to codegen_output=), so you can inspect — and even edit — the code that will actually be compiled.

To support editing the generated code and re-running with your edits applied, IR Lower Trace maintains three cooperating files:

File	Role
codegen.cpp	Working copy — user-editable. This is what gets compiled when you rerun.
codegen.cpp.original	Baseline — the codegen snapshot the working copy was last synced from. Written only on init or re-sync, never blindly overwritten.
codegen.cpp.latest	Latest codegen output — the actual output of the most recent run, overwritten every run for diff reference.

On each run a three-way comparison (baseline / working copy / current codegen output) decides how to proceed:

Situation	codegen.cpp vs .original	.latest vs .original	Action	Console tag
No change	identical	identical	Compile with codegen output as-is	—
Codegen changed only	identical	differs	Regenerate codegen.cpp and .original from new codegen	REGENERATED
User edited only	differs	identical	Inject the working copy (PATCHED)	PATCHED
Both changed, working == latest	differs	differs (working matches latest)	Advance baseline; use working copy	SYNCED
Both changed, working != latest	differs	differs (working differs from latest)	CONFLICT — back up working copy → .bak (conflict backup) and old baseline → .original.bak, then regenerate from new codegen	CONFLICT
First run (no baseline)	—	—	Initialise .original, copy to codegen.cpp	(init)
codegen.cpp exists without baseline	—	—	Back up pre-existing codegen.cpp → .bak (safety backup), then initialise baseline	INIT-BACKUP

Note on .bak files: The backups created by CONFLICT and INIT-BACKUP serve different purposes. INIT-BACKUP preserves a pre-existing codegen.cpp of unknown origin before the trace tool takes it over. CONFLICT preserves the user’s edits before a codegen change overwrites them. Recover CONFLICT edits with diff codegen.cpp.original.bak codegen.cpp.bak.

Typical Workflow

1. Inspect — Run once with TL_LOWER_TRACE=1. Open codegen.cpp to read the generated source.

2. Edit — Modify codegen.cpp (e.g. add a printf, tweak a loop). Do not touch .original.

3. Rerun — Run again. Because codegen.cpp differs from .original but codegen output is unchanged, you’ll see PATCHED from …/codegen.cpp and your edited source is compiled.

4. Iterate — Keep editing and rerunning. Each run re-injects your working copy.

5. If codegen itself changes (e.g. you modified the TileLang program) — two outcomes:

   • If your edits happen to match the new codegen output → SYNCED (baseline advances, your edits preserved).

   • If both your edits and codegen changed and they differ → CONFLICT. Your working copy is backed up to codegen.cpp.bak and the old baseline to codegen.cpp.original.bak. Recover your edits with diff codegen.cpp.original.bak codegen.cpp.bak, then re-apply them against the freshly regenerated codegen.cpp.

Note

Backend requirements for edit-and-recompile. The edit-and-recompile workflow requires a source-compiling execution backend — nvrtc, cython, or cutedsl. These backends use *_without_compile codegen FFIs that produce source-only modules, then compile the (edited) source string at runtime via NVRTC / Cython / CuTeDSL.

The default tvm_ffi backend pre-compiles device code to a binary (PTX/hsaco) from TIR during codegen. When the tvm_ffi backend is active and you edit codegen.cpp, you’ll see a NOTE message indicating that your edits are recorded in the trace for diff viewing but were not recompiled. To use edit-and-recompile, switch to a source-compiling backend:

# For CUDA targets:
tilelang.compile(..., execution_backend="nvrtc")

# For HIP targets:
tilelang.compile(..., execution_backend="cython")

How It Works

IR Lower Trace installs three layers of transparent hooks (all via monkey-patch, restored by disable()):

1. tvm.ir.transform.Pass.__call__ — every pass invocation is intercepted to capture str(mod) before and after, compute +/− line counts, and append a LowerRecord. Passes that run outside any pipeline window are tagged with the unscoped phase.

2. PassPipeline.lower (new architecture) or phase functions (legacy architecture) — sets the current phase context so passes invoked within a pipeline run are grouped under a label like pipeline_c. Legacy phase functions are discovered via AST scanning (_discover_passes) and bytecode inspection.

3. Codegen FFI (target.build.tilelang_cuda, …_hip, …_c, …_llvm, etc.) — captures the final TIR → source lowering and drives the three-file edit-recompile workflow described above.

Pass records are appended at runtime (not pre-registered), so conditional passes that are skipped at runtime — e.g. LetInline when should_force_let_inline() is False — simply do not appear, leaving no phantom/skipped slots. The HTML report is flushed incrementally after every pass (O(n) total cost), so partial results survive even a crash or SIGKILL.

When the same process compiles multiple kernels, each PassPipeline.lower invocation increments a run counter and tags phases with a run2_, run3_, … prefix; all records accumulate into a single report so you can compare runs side by side.

Note

The pipeline hook wraps tvm.ir.transform.Pass.__call__ for the entire process, so it observes TVM passes outside the immediate tilelang.compile() call as well. Capturing IR and, in HTML mode, rewriting the report after every pass adds debugging overhead. Leave the feature disabled for normal builds and benchmarks.

Choosing an Output Mode

• Use terminal for a small number of passes and immediate feedback.

• Use html when a lowering pipeline produces many changes or you need to compare complete before and after scripts.

• Use both when iterating interactively but retaining a shareable report.

IR Lower Trace compares the textual form returned by IRModule.script(). A textual change is useful evidence about a pass, but it does not by itself establish a semantic or performance change.

Tips

• Use terminal mode for quick checks — the colored diff prints as passes run, so you can see changes in real time.

• Use html mode for thorough analysis — navigate across many passes, expand hidden context, and copy IR snippets.

• Combine with TL_LOWER_TRACE_DIR to direct reports to a specific location, e.g. when running in CI or comparing across runs.

• The hook captures all passes in the lowering pipeline, including those triggered internally by tilelang.compile(). This makes it useful for understanding the full compilation flow.

• If you previously used TILELANG_PASS_DIFF, switch to TL_LOWER_TRACE — it is a strict superset and is the tool that receives future improvements.