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).

• Per-compilation isolation — each kernel compilation owns its records, pass indices, pipeline scopes, and report lifecycle. Concurrent kernels do not write into one shared in-memory trace.

• 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 registers a tool factory for future compilation sessions. It does not patch TVM’s default PassContext or global codegen functions. To change the setting after import, call enable() programmatically.

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_device.cpp              # generated device source (editable, see below)
    ├── codegen_device.cpp.original     # device baseline for edit/recompile workflow
    ├── codegen_device.cpp.latest       # latest generated device source
    ├── codegen_host.cpp                # generated CPU wrapper, when host codegen runs
    ├── codegen_host.cpp.original       # independent host baseline
    ├── codegen_host.cpp.latest         # latest generated host source
    └── .run_records/
        └── run_<YYYYMMDD_HHMMSS_ffffff>_<pid>_<unique-id>/
            ├── 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_device_before.tir
            │   ├── 42_codegen_device_after.cpp
            │   ├── 49_codegen_host_before.tir
            │   └── 49_codegen_host_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. Both device and host records belong to the codegen phase, with names codegen_device and codegen_host respectively. Host files and records appear only when the execution backend invokes host codegen. Each working file has its own baseline, latest output, and backups; host codegen cannot overwrite the device working file or its editing state.

The standard backend wrappers explicitly identify each event as device or host codegen. Generic instrumentation calls that omit role leave it unspecified (None): these retain the codegen record name and default codegen.cpp working file. An explicit codegen_output path is used unchanged for such events. Custom callers should supply a role to distinguish device and host output; unspecified events share the legacy working file and edit state.

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.

Compilation instrumentation: enable() / disable()

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	Device working-file path; a host working-file path is derived when host codegen runs. Pass None explicitly to suppress both working files.	<script_dir>/codegen_device.cpp and <script_dir>/codegen_host.cpp in html/both mode; no working files in terminal mode

An explicit codegen_output="foo.cpp" writes device source to foo.cpp and host source to foo_host.cpp. The host path inserts _host before the final extension: foo.kernel.cpp becomes foo.kernel_host.cpp, and a path without an extension, such as foo, becomes foo_host. Each path has its own .original, .latest, and backup files. Setting codegen_output=None suppresses these stable working files and their edit state, while codegen records still appear in the trace and raw snapshots remain available.

enable() is idempotent — calling it multiple times is safe. Its values are snapshotted when a compilation begins, so changing them cannot mutate a kernel that is already compiling.

When to use disable()

Function	When to call	What it does
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)	Unregisters the tool factory. A compilation already in progress keeps its own session and completes normally.

from tilelang.tools import lower_trace as lt

lt.enable(mode="both")

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

# Second kernel — traced into its own session and report automatically.
kernel2 = tilelang.compile(func_b)

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

Note

If lt.disable() is not called, tracing stays active for future compilations. Each HTML report is finalized when its owning compilation exits, including when that compilation raises.

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, device and host codegen steps (TIR → C/CUDA/HIP/…) are intercepted separately. Device source runs on the GPU or NPU and contains the computation, device memory transfers, and synchronization. Host source runs on the CPU and wraps runtime operations such as argument checks, output allocation, and launching the device kernel.

In html/both mode, the generated sources are written to <script_dir>/codegen_device.cpp and <script_dir>/codegen_host.cpp by default. For CUDA and Ascend, tvm_ffi invokes both device and host codegen; nvrtc and cython do not invoke a separate host codegen step at this layer. Their runtime adapters still provide CPU-side launch logic. Only backends that invoke host codegen produce codegen_host.cpp.

To support editing the generated code and re-running with your edits applied, IR Lower Trace maintains three cooperating files for each codegen role. The device files are shown below; the host files use the same suffixes on codegen_host.cpp and keep an independent edit state:

File	Role
codegen_device.cpp	Working copy — user-editable. Edits are recompiled when the backend supports the workflow described below.
codegen_device.cpp.original	Baseline — the codegen snapshot the working copy was last synced from. Written only on init or re-sync, never blindly overwritten.
codegen_device.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	Working copy vs .original	.latest vs .original	Action	Console tag
No change	identical	identical	Compile with codegen output as-is	—
Codegen changed only	identical	differs	Regenerate the working copy 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 the working file	(init)
Working copy exists without baseline	—	—	Back up the pre-existing working file → .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 working file 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_device.cpp.original.bak codegen_device.cpp.bak for the device working copy, or the equivalent command with codegen_host.cpp for the host working copy. These backup files are created only when needed.

Typical Workflow

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

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

3. Rerun — Run again. Because codegen_device.cpp differs from .original but codegen output is unchanged, you’ll see PATCHED from …/codegen_device.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_device.cpp.bak and the old baseline to codegen_device.cpp.original.bak. Recover your edits with diff codegen_device.cpp.original.bak codegen_device.cpp.bak, then re-apply them against the freshly regenerated codegen_device.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_device.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")

Capturing host source and preserving host edits does not add host edit-and-recompile support; the existing backend requirements still apply.

How It Works

Instrumentation sessions are owned by complete compiler entry points such as tilelang.compile(), tilelang.lower(), and grouped compilation. Lower-level pipeline and codegen components contribute events to the active session but do not create standalone sessions themselves.

IR Lower Trace uses three cooperating layers:

1. TVM PassInstrument — before/after callbacks capture str(mod), compute +/− line counts, and append a LowerRecord. The shared nested callback stack records depth and parent relationships, including passes invoked internally by another TVM pass. Passes outside a pipeline window are tagged with the unscoped phase.

2. Explicit PassPipeline.lower scope — contributes a context-local phase to the caller-owned session so passes in a backend pipeline are grouped under labels such as pipeline_c. The pipeline does not create or finalize the compile session. This is a normal pipeline boundary, not a class monkey-patch.

3. Explicit backend codegen hook — TileLang’s device/host codegen registry routes each target.build.* call through the active compile session. This captures final TIR → source lowering, labels events as device or host, and drives each role’s three-file edit-recompile workflow without replacing a process-global FFI. When no session is active, the registry invokes the underlying codegen directly.

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 one logical compilation uses multiple PassContexts (for example grouped lowering followed by device and host codegen), they obtain fresh callback objects from the same compile session. Their records share one monotonic index inside that session. A different compilation starts from index zero with a new record list.

Note

enable() affects TileLang-managed compilation sessions only; it never mutates the PassContext that happens to be current at enable time. For a standalone TVM PassContext, lt.create_pass_instrument() remains available for manual use. 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.