tilelang.instrumentation.session¶

Compile-scoped orchestration for TileLang instrumentation tools.

Classes¶

CodegenEvent

One explicit backend code-generation call within a compile session.

PassInstrumentationTool

Per-compile extension point for pass and codegen developer tools.

CompilePassInstrumentationSession

Concrete pass-tool state owned by one logical compilation.

Functions¶

register_pass_instrumentation_tool(name, factory)

Register or replace a factory used by future compile sessions.

unregister_pass_instrumentation_tool(name)

Stop adding a tool to future compile sessions.

current_compile_pass_instrumentation()

Return the compile session active in this execution context.

current_pass_instrument_context()

Return metadata for the PassContext whose instruments are being created.

compile_pass_instrumentation([name, tools, ...])

Own pass-tool state for one logical compilation.

create_pass_instruments(*[, context])

Create instruments for the active compile session, if any.

instrument_current_pass_context()

Temporarily add active-session instruments to TVM's current context.

run_codegen_with_instrumentation(name, mod, target, call)

Run an explicit backend codegen call through the active session.

pass_pipeline(name)

Run a backend pipeline under its phase and compile-session tools.

Module Contents¶

class tilelang.instrumentation.session.CodegenEvent¶

One explicit backend code-generation call within a compile session.

name: str¶
mod: Any¶
target: Any¶
class tilelang.instrumentation.session.PassInstrumentationTool¶

Per-compile extension point for pass and codegen developer tools.

pass_instrument_priority = 0¶

Ordering key for PassContext callbacks; lower values run first.

create_pass_instrument()¶

Create fresh callback state for one TVM PassContext.

Return type:

object | None

pipeline_scope(base_phase)¶

Return the scope entered around one backend pass pipeline.

Parameters:

base_phase (str)

Return type:

contextlib.AbstractContextManager[None]

run_codegen(event, next_call)¶

Observe or wrap codegen, delegating to next_call exactly once.

Parameters:
  • event (CodegenEvent)

  • next_call (collections.abc.Callable[[], Any])

Return type:

Any

finish(error)¶

Finalize this tool after its owning compile session exits.

Parameters:

error (BaseException | None)

Return type:

None

class tilelang.instrumentation.session.CompilePassInstrumentationSession(tools, *, name=None)¶

Concrete pass-tool state owned by one logical compilation.

Parameters:
name = None¶
tools¶
create_pass_instruments(*, context=None)¶

Create independent, ordered instruments for one TVM PassContext.

Parameters:

context (str | None)

Return type:

list[object]

pipeline_scope(base_phase)¶

Enter every tool scope for one backend pass pipeline.

Parameters:

base_phase (str)

Return type:

collections.abc.Generator[None, None, None]

run_codegen(event, call)¶

Compose tool codegen middleware in registration order.

Parameters:
  • event (CodegenEvent)

  • call (collections.abc.Callable[[], Any])

Return type:

Any

finish(error)¶

Finalize tools in reverse registration order.

Parameters:

error (BaseException | None)

Return type:

None

find_tool(tool_type)¶

Return the first tool matching tool_type, if present.

Parameters:

tool_type (type[_ToolT])

Return type:

_ToolT | None

tilelang.instrumentation.session.register_pass_instrumentation_tool(name, factory)¶

Register or replace a factory used by future compile sessions.

Parameters:
Return type:

None

tilelang.instrumentation.session.unregister_pass_instrumentation_tool(name)¶

Stop adding a tool to future compile sessions.

Parameters:

name (str)

Return type:

None

tilelang.instrumentation.session.current_compile_pass_instrumentation()¶

Return the compile session active in this execution context.

Return type:

CompilePassInstrumentationSession | None

tilelang.instrumentation.session.current_pass_instrument_context()¶

Return metadata for the PassContext whose instruments are being created.

Return type:

str | None

tilelang.instrumentation.session.compile_pass_instrumentation(name=None, *, tools=(), include_default_tools=True, reuse_existing=True)¶

Own pass-tool state for one logical compilation.

Nested compiler helpers reuse the active session by default, so every PassContext involved in one compilation writes to the same tool state. Independent tools such as Pass Visualizer can request a fresh session by setting reuse_existing=False and include_default_tools=False.

Parameters:
Return type:

collections.abc.Generator[CompilePassInstrumentationSession, None, None]

tilelang.instrumentation.session.create_pass_instruments(*, context=None)¶

Create instruments for the active compile session, if any.

context is descriptive metadata for this particular TVM PassContext. Tools can snapshot it during callback creation without changing their create_pass_instrument interface.

Parameters:

context (str | None)

Return type:

list[object]

tilelang.instrumentation.session.instrument_current_pass_context()¶

Temporarily add active-session instruments to TVM’s current context.

Compiler front doors use this only when they created the compile session themselves. Nested helpers normally run inside a caller-owned PassContext whose instruments were created from the same session, so they must not attach a duplicate set.

Return type:

collections.abc.Generator[None, None, None]

tilelang.instrumentation.session.run_codegen_with_instrumentation(name, mod, target, call)¶

Run an explicit backend codegen call through the active session.

Parameters:
  • name (str)

  • mod (Any)

  • target (Any)

  • call (collections.abc.Callable[[], _CodegenResult])

Return type:

_CodegenResult

tilelang.instrumentation.session.pass_pipeline(name)¶

Run a backend pipeline under its phase and compile-session tools.

Parameters:

name (str)

Return type:

collections.abc.Generator[None, None, None]