tilelang.language.eager.builder¶

Attributes¶

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

類別¶

`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`
`Macro`
`TirTemplate`	Template for generating TIR PrimFunc with dynamic shape substitution.
`JITFunc`	Internal wrapper for JIT-compiled 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

回傳型別:: 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)¶

參數:: v (Any)
回傳型別:: bool

tilelang.language.eager.builder.EagerJITStage¶

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: EagerJITStage = 'none'¶

eager_jit_subs: dict[str, tvm.tir.expr.PrimExpr]¶

current_file = '<unknown>'¶

current_line = 0¶

current_macro_name = '<unknown-macro>'¶

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

classmethod current()¶

回傳型別:: Self

prim_func(name)¶

macro(name=None, annotations=None)¶

get()¶

回傳型別:: PrimFunc

find_frame_idx(frame, start=0)¶

參數:: frame (type | tuple[type, Ellipsis])
回傳型別:: int | None

enter_frame(frame)¶

參數:: frame (contextlib.AbstractContextManager[Any])

check_continue_break()¶

with_frame(frame)¶

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

ctx_if(cond)¶

ctx_then(val)¶

ctx_else(val)¶

eval(val)¶

參數:: 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)¶

參數:

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

aug_assign(op, target, aug_value, name=None)¶

參數:: name (str | None)

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

參數:

name (str)
value (Any)

回傳型別:

Any

macro_arg(name, value)¶

prim_func_arg(name, value)¶

arg(name, value)¶

override(name)¶

參數:: name (str)

constexpr(name, dtype='int32')¶

參數:

name (str)
dtype (str)

回傳型別:

tvm.tir.expr.Var

set_fileline(filename, lineno, name)¶

參數:

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

get_fileline_stack(stacklevel=1)¶

skip_kernel_ctx()¶

class tilelang.language.eager.builder.PrimFunc¶

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

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]

name: str¶

orig_func: Callable[_P, _T]¶

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

annotations: dict[str, Any]¶

property source: str¶

回傳型別:: str

__call__(*args, **kwargs)¶

參數:

args (_P)
kwargs (_P)

回傳型別:

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

回傳:

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

參數:

func (Callable[_P, _T])

回傳型別:

Macro[_P, _T]

也參考

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

參數:

name (str)
dtype (str)

回傳型別:

tvm.tir.expr.Var | 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.

name: str¶

prim_func: PrimFunc[_P, _T]¶

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

constexprs: set[tvm.tir.expr.Var] = None¶

is_lazy_style: bool = False¶

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

classmethod create(name, prim_func, constexpr, ir_gen=None)¶

參數:

name (str)
prim_func (PrimFunc[_P, _T])
constexpr (set[tvm.tir.expr.Var])
ir_gen (tilelang.language.eager.ast.IRGenerator[_P, _T] | None)

回傳型別:

TirTemplate[_P, _T]

classmethod from_lazy_style(name, prim_func)¶

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

參數:

name (str)
prim_func (PrimFunc[_P, _T])

回傳型別:

TirTemplate[_P, _T]

get_tir(tensor_args, given_tensor_args, 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).

參數:: 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)¶

參數:

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

回傳型別:

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