tilelang.tools.pass_visualizer.core¶

Core helpers for the pass visualizer: load a user TileLang kernel, capture real compiler-pass executions, and render a PrimFunc’s SBlock structure tree.

StructureTreePassInstrument observes the canonical CUDA lowering prologue through TVM’s PassInstrument API. The visualizer therefore follows the passes that actually execute instead of maintaining a duplicate pass list.

The kernel file is taken as input; any @tilelang.jit kernel in the file is auto-discovered. These helpers are consumed by viewer.py to emit an interactive HTML pass browser.

Attributes¶

logger

Classes¶

`PassStructureRecord`	Before/after structure snapshots for one top-level compiler pass.
`StructureTreePassInstrument`	Capture structure trees around the real top-level passes in a pipeline.
`StructureTreePassTool`	Per-viewer tool that creates its PassContext-local capture instrument.

Functions¶

`load_user_module`(path)	Import an arbitrary user TileLang source file as a module.
`discover_jit_kernels`(module)	Find every @tilelang.jit object (JITImpl) defined at module top level.
`kernel_to_tir`(kernel, **kwargs)	Elaborate a kernel into its un-lowered PrimFunc (TIR).
`build_module`(func[, target])	Wrap a PrimFunc into an IRModule and resolve the (target, target_host) pair.
`inspect_structure`(mod)	Print each PrimFunc top-down: params → buffer_map → attrs → body (SBlock tree).
`capture_structure`(mod)	Render `inspect_structure` into stable text lines for diffing.

Module Contents¶

tilelang.tools.pass_visualizer.core.logger¶

tilelang.tools.pass_visualizer.core.load_user_module(path)¶

Import an arbitrary user TileLang source file as a module.

Parameters:: path (str)

tilelang.tools.pass_visualizer.core.discover_jit_kernels(module)¶

Find every @tilelang.jit object (JITImpl) defined at module top level.

Return type:: dict[str, tilelang.jit.JITImpl]

tilelang.tools.pass_visualizer.core.kernel_to_tir(kernel, **kwargs)¶

Elaborate a kernel into its un-lowered PrimFunc (TIR).

Accepts three forms:

JITImpl (@tilelang.jit) -> .get_tir(**kwargs)
a factory callable (@T.prim_func wrapper returning a PrimFunc) -> call it
a PrimFunc already -> returned as-is

Return type:: tvm.tirx.PrimFunc

tilelang.tools.pass_visualizer.core.build_module(func, target='auto')¶

Wrap a PrimFunc into an IRModule and resolve the (target, target_host) pair.

Parameters:

func (tvm.tirx.PrimFunc)
target (str | tvm.target.Target)

tilelang.tools.pass_visualizer.core.inspect_structure(mod)¶

Print each PrimFunc top-down: params → buffer_map → attrs → body (SBlock tree).

Parameters:: mod (tilelang.tvm.IRModule)
Return type:: None

tilelang.tools.pass_visualizer.core.capture_structure(mod)¶

Render inspect_structure into stable text lines for diffing.

Parameters:: mod (tilelang.tvm.IRModule)
Return type:: list[str]

class tilelang.tools.pass_visualizer.core.PassStructureRecord¶

Before/after structure snapshots for one top-level compiler pass.

name: str¶

sequence: int¶

before_lines: list[str]¶

after_lines: list[str]¶

property changed: bool¶

Return type:: bool

class tilelang.tools.pass_visualizer.core.StructureTreePassInstrument¶

Bases: tilelang.instrumentation.StackedPassInstrument

Capture structure trees around the real top-level passes in a pipeline.

Some top-level TileLang passes invoke nested TVM passes internally. The shared stack instrument tracks every callback for pairing, while this consumer snapshots only depth-zero passes so the browser remains linear.

property records: list[PassStructureRecord]¶

Return type:: list[PassStructureRecord]

property input_lines: list[str] | None¶

Return type:: list[str] | None

property incomplete_passes: list[str]¶

Return type:: list[str]

ordered_records()¶

Return completed top-level pass records in execution order.

Return type:: list[PassStructureRecord]

class tilelang.tools.pass_visualizer.core.StructureTreePassTool¶

Bases: tilelang.instrumentation.PassInstrumentationTool

Per-viewer tool that creates its PassContext-local capture instrument.

instrument: StructureTreePassInstrument | None = None¶

create_pass_instrument()¶

Create fresh callback state for one TVM PassContext.

Return type:: StructureTreePassInstrument