tilelang.tools.pass_visualizer.core¶

Core helpers for the pass visualizer: load a user TileLang kernel, build the CUDA lowering pass pipeline, and render a PrimFunc’s SBlock structure tree.

This mirrors CUDAPassPipelineBodyPrologue in tilelang/cuda/pipeline.py, running each prologue pass (through LayoutInference and the post-LayoutInference lowering passes) so the structure tree can be captured before/after every stage.

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.

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.

build_pass_stages(target)

Build the ordered CUDA prologue pass list, through the full Prologue.

inspect_structure(mod)

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

Module Contents¶

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.build_pass_stages(target)¶

Build the ordered CUDA prologue pass list, through the full Prologue.

Each stage is a (name, transform) pair where transform is a TVM pass object callable as transform(mod) -> mod. Conditional passes are resolved here so the returned list reflects what actually runs for this target/config.

Matches CUDAPassPipelineBodyPrologue (tilelang/cuda/pipeline.py:68-137).

Parameters:

target (tvm.target.Target)

Return type:

list[tuple[str, object]]

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