tilelang.language.eager.builder¶

Attributes¶

`logger`
`thread_local_storage`
`ContinueOrBreak`
`AnyFrame`
`TIR_CONTROL_FRAME`
`TIR_VAR_SCOPE_FRAME`

Classes¶

`Frame`	Frame are virtual context managers used in frontend only
`MacroFrame`	Frame are virtual context managers used in frontend only
`ExitedMacroFrame`	Frame are virtual context managers used in frontend only
`BoolOpFrame`	Frame are virtual context managers used in frontend only
`ContinueFrame`	Frame are virtual context managers used in frontend only
`BreakFrame`	Frame are virtual context managers used in frontend only
`SerialForWithStep`
`OutTensor`
`Ref`
`UnrollForWithStep`
`Builder`
`PrimFunc`	Abstract base class for generic types.
`Macro`	Abstract base class for generic types.
`TirTemplate`	Template for generating TIR PrimFunc with dynamic shape substitution.
`JITFunc`	Internal wrapper for JIT-compiled functions.

Functions¶

`unwrap_expr`(expr)	unwrap expr and convert it into PrimExpr like
`unwrap_cond`(expr)	unwrap expr and convert to bool condition
`is_var`(v)
`macro`([func])	Decorator that converts a Python function into a TileLang macro.
`get_type_hints`(func)
`const`(name[, dtype])	Declare constexpr variables for dynamic tensor dimensions (eager mode only).
`substitute_primfunc`(prim_func, vmap)
`prim_func`([func, eager_jit])

Module Contents¶

tilelang.language.eager.builder.logger¶

tilelang.language.eager.builder.unwrap_expr(expr)¶

unwrap expr and convert it into PrimExpr like

Return type:: tvm.tir.expr.PrimExpr | int | float

tilelang.language.eager.builder.unwrap_cond(expr)¶: unwrap expr and convert to bool condition

tilelang.language.eager.builder.thread_local_storage¶

class tilelang.language.eager.builder.Frame¶

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

__enter__()¶

__exit__(exc_type, exc_value, traceback)¶

class tilelang.language.eager.builder.MacroFrame¶

Bases: Frame

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

class tilelang.language.eager.builder.ExitedMacroFrame¶

Bases: Frame

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

class tilelang.language.eager.builder.BoolOpFrame¶

Bases: Frame

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

class tilelang.language.eager.builder.ContinueFrame¶

Bases: Frame

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

class tilelang.language.eager.builder.BreakFrame¶

Bases: Frame

Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR.

class tilelang.language.eager.builder.SerialForWithStep¶

start: tvm.tir.expr.PrimExpr¶

stop: tvm.tir.expr.PrimExpr¶

step: tvm.tir.expr.PrimExpr¶

annotations: dict[str, Any] | None = None¶

class tilelang.language.eager.builder.OutTensor¶

shape: collections.abc.Sequence[tvm.tir.expr.PrimExpr]¶

dtype: tilelang.language.dtypes.dtype¶

property strides¶

class tilelang.language.eager.builder.Ref¶

bufload: tvm.tir.expr.BufferLoad¶

property buffer¶

store(value)¶

load()¶

class tilelang.language.eager.builder.UnrollForWithStep¶: Bases: SerialForWithStep

tilelang.language.eager.builder.ContinueOrBreak¶

tilelang.language.eager.builder.AnyFrame¶

tilelang.language.eager.builder.TIR_CONTROL_FRAME¶

tilelang.language.eager.builder.TIR_VAR_SCOPE_FRAME¶

tilelang.language.eager.builder.is_var(v)¶

Parameters:: v (Any)
Return type:: bool

class tilelang.language.eager.builder.Builder¶

Bases: tilelang.language.eager.ast.BaseBuilder

frames: list[AnyFrame] = []¶

ir_builder¶

name_inside_frame: dict[str, AnyFrame]¶

macro_arg_annot¶

out_idx = []¶

out_tensor_cnt = 0¶

constexpr_var¶

eager_jit = False¶

current_file = '<unknown>'¶

current_line = 0¶

current_macro_name = '<unknown-macro>'¶

macro_fileline_stack: list[tuple[str, int, str]] = []¶

classmethod current()¶

Return type:: Self

prim_func(name)¶

macro(name=None, annotations=None)¶

get()¶

Return type:: PrimFunc

find_frame_idx(frame, start=0)¶

Parameters:: frame (type | tuple[type, Ellipsis])
Return type:: int | None

enter_frame(frame)¶

Parameters:: frame (contextlib.AbstractContextManager[Any])

check_continue_break()¶

with_frame(frame)¶

Parameters:: frame (contextlib.AbstractContextManager[Any] | None)

ctx_if(cond)¶

ctx_then(val)¶

ctx_else(val)¶

eval(val)¶

Parameters:: val (Any)

ctx_for(it)¶

ctx_continue()¶

ctx_break()¶

ctx_while(cond)¶

bind(name, value, annot=BaseBuilder.empty)¶

unwrap_value(value)¶: Unwrap some tilelang objects to get their inner value

bind_immutable(name, value)¶: Bind an immutable tilelang objects. The immutability means the result is usually not changed or re-assigned in a python block.

assign_slice(lval, sl, value, annot=BaseBuilder.empty)¶

Parameters:

lval (Any)
sl (slice)
value (Any)

aug_assign(op, target, aug_value)¶

aug_assign_slice(op, target, sl, aug_value)¶

boolop(op, left, right=None)¶

ifexp(cond, then, otherwise)¶

ret(value=None)¶

ctx_with(ctx)¶

assert_expr(cond, msg=None)¶

rval(name, value)¶

Parameters:

name (str)
value (Any)

Return type:

Any

macro_arg(name, value)¶

prim_func_arg(name, value)¶

arg(name, value)¶

override(name)¶

Parameters:: name (str)

constexpr(name, dtype='int32')¶

Parameters:

name (str)
dtype (str)

Return type:

tvm.tir.expr.Var

set_fileline(filename, lineno, name)¶

Parameters:

filename (str)
lineno (int)
name (str)

get_fileline_stack(stacklevel=1)¶

class tilelang.language.eager.builder.PrimFunc¶

Bases: Generic[_P, _T], tvm.tir.PrimFunc

Abstract base class for generic types.

A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as:

class Mapping(Generic[KT, VT]):
    def __getitem__(self, key: KT) -> VT:
        ...
    # Etc.

This class can then be used as follows:

def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
    try:
        return mapping[key]
    except KeyError:
        return default

params: list[tvm.tir.Var | tvm.tir.Buffer]¶

body: tvm.tir.Stmt¶

ret_type: tvm.ir.Type¶

buffer_map: tvm_ffi.container.Map[tvm.tir.Var, tvm.tir.Buffer]¶

attrs: tvm.Attrs | None¶

span: tvm.ir.base.Span | None¶

ir_gen: tilelang.language.eager.ast.IRGenerator[_P, _T] | None¶

orig_func: Callable[_P, _T] | None¶

out_idx_override: list[int] | None¶

class tilelang.language.eager.builder.Macro¶

Bases: Generic[_P, _T]

Abstract base class for generic types.

A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as:

class Mapping(Generic[KT, VT]):
    def __getitem__(self, key: KT) -> VT:
        ...
    # Etc.

This class can then be used as follows:

def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
    try:
        return mapping[key]
    except KeyError:
        return default

name: str¶

orig_func: Callable[_P, _T]¶

ir_gen: tilelang.language.eager.ast.IRGenerator[_P, _T]¶

annotations: dict[str, Any]¶

property source: str¶

Return type:: str

__call__(*args, **kwargs)¶

Parameters:

args (_P)
kwargs (_P)

Return type:

_T

__hash__()¶

__eq__(other)¶

tilelang.language.eager.builder.macro(func=None)¶

Decorator that converts a Python function into a TileLang macro. TileLang macro is very similar to PrimFunc, it can be used in prim_func or another macro. :param func: The Python function to be converted into a macro. This function will be analyzed

and transformed into an IR generation function. The function can take any parameters (_P) and return any type (_T).

Returns:

Macro[_P, _T] – A Macro object that wraps the original function with IR generation capabilities. The returned Macro preserves the original function’s signature (parameters _P and return type _T) while adding metaprogramming capabilities.
Example
——– – >>> @macro … def my_macro(x: T.int32) -> T.int32: … return x ** 2 >>> @prim_func … def my_func(A: T.Tensor((10,), T.int32), B: T.Tensor((10,), T.int32)): … with T.Kernel(1) as _: … for i in T.serial(10): … B[i] = my_macro(A[i])

Parameters:

func (Callable[_P, _T])

Return type:

Macro[_P, _T]

See also

Macro: The class that wraps macro functions
mutate: The function that transforms Python code into IR generators

tilelang.language.eager.builder.get_type_hints(func)¶

tilelang.language.eager.builder.const(name, dtype='int32')¶

Declare constexpr variables for dynamic tensor dimensions (eager mode only).

In eager mode, use T.const() to declare shape dimensions that will be inferred from actual tensor arguments at runtime.

Example:

@tilelang.jit
def kernel(A, B):
    M, N = T.const("M, N")
    A: T.Tensor[[M, N], T.float32]
    ...

Parameters:

name (str)
dtype (str)

Return type:

tuple[tvm.tir.expr.Var, Ellipsis]

class tilelang.language.eager.builder.TirTemplate¶

Bases: Generic[_P, _T]

Template for generating TIR PrimFunc with dynamic shape substitution.

For lazy-style functions, the PrimFunc is used directly without substitution. For eager-style functions, constexpr variables are substituted based on actual tensor shapes at runtime.

prim_func: PrimFunc[_P, _T]¶

matcher: dict[tvm.tir.expr.Var, tuple[tvm.tir.Var, str, int]] | None = None¶

is_lazy_style: bool = False¶

classmethod create(prim_func, constexpr)¶

Parameters:

prim_func (PrimFunc[_P, _T])
constexpr (set[tvm.tir.expr.Var])

Return type:

TirTemplate[_P, _T]

classmethod from_lazy_style(prim_func)¶

Create template from lazy-style function that returns PrimFunc directly.

Parameters:: prim_func (PrimFunc[_P, _T])
Return type:: TirTemplate[_P, _T]

get_tir(**kwargs)¶

class tilelang.language.eager.builder.JITFunc¶

Bases: Generic[_P, _T]

Internal wrapper for JIT-compiled functions.

This class handles both lazy and eager execution styles:

lazy style: Function explicitly returns a PrimFunc. The original function is called directly to obtain the TIR.
eager style: Function uses the DSL builder pattern with tensor type annotations. The TIR is constructed by tracing the function body through the Builder.

The style is determined by _is_lazy_style() which checks if calling the original function returns a PrimFunc directly.

orig_func: Callable[_P, _T]¶

arg_names: list[str]¶

tensor_args: dict[str, tvm.tir.Buffer | tvm.tir.expr.Var]¶

tensor_args_defaults: dict[str, Any]¶

ir_gen: tilelang.language.eager.ast.IRGenerator[_P, _T]¶

mode: Literal['auto', 'lazy', 'eager'] = 'auto'¶

__post_init__()¶

parse_args(*args, **kwargs)¶: Parse arguments and return cache key and tensor args.

get_tir(*args, **kwargs)¶

__call__(*args, **kwargs)¶

set_mode(mode)¶

Set the JIT execution mode (internal use only).

Parameters:: mode (Literal['lazy', 'eager'])

__getattr__(name)¶

tilelang.language.eager.builder.substitute_primfunc(prim_func, vmap)¶

tilelang.language.eager.builder.prim_func(func=None, *, eager_jit=False)¶

Parameters:

func (Callable[_P, _T])
eager_jit (bool)

Return type:

PrimFunc[_P, _T] | JITFunc[_P, _T]