tilelang.language.v2.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
`ConstIfFrame`	Frame are virtual context managers used in frontend only
`BlockFrame`	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`
`PrimFuncCreater`
`PrimFunc`
`Macro`

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)
`prim_func`([func, generator])	Decorator to create a primitive function (PrimFunc) for TileLang IR generation.

Module Contents¶

tilelang.language.v2.builder.logger¶

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

unwrap expr and convert it into PrimExpr like

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

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

tilelang.language.v2.builder.thread_local_storage¶

class tilelang.language.v2.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.v2.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.v2.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.v2.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.v2.builder.ConstIfFrame¶

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.v2.builder.BlockFrame¶

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.v2.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.v2.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.v2.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.v2.builder.OutTensor¶

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

dtype: tilelang.language.v2.dtypes.dtype¶

property strides¶

class tilelang.language.v2.builder.Ref¶

bufload: tvm.tir.expr.BufferLoad¶

property buffer¶

store(value)¶

load()¶

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

tilelang.language.v2.builder.ContinueOrBreak¶

tilelang.language.v2.builder.AnyFrame¶

tilelang.language.v2.builder.TIR_CONTROL_FRAME¶

tilelang.language.v2.builder.TIR_VAR_SCOPE_FRAME¶

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

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

class tilelang.language.v2.builder.Builder(func_annot=None)¶

Bases: tilelang.language.v2.ast.BaseBuilder

Parameters:: func_annot (tilelang.language.v2.annot.FuncAnnot)

frames: list[AnyFrame] = []¶

ir_builder¶

name_inside_frame: dict[str, AnyFrame]¶

macro_arg_annot¶

func_annot = None¶

out_idx = []¶

out_tensor_cnt = 0¶

arg_vt¶

classmethod current()¶

Return type:: Self

prim_func(name)¶

macro(name=None, annotations=None)¶

get()¶

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)

class tilelang.language.v2.builder.PrimFuncCreater¶

Bases: Generic[_P, _T]

func_annot: tilelang.language.v2.annot.FuncAnnot¶

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

orig_func: Callable[_P, _T]¶

property annot: dict[str, tilelang.language.v2.annot.Annot]¶

Return type:: dict[str, tilelang.language.v2.annot.Annot]

__call__(*args, **kwargs)¶

Parameters:

args (_P)
kwargs (_P)

Return type:

PrimFunc[_P, _T]

__repr__()¶

class tilelang.language.v2.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.v2.ast.IRGenerator[_P, _T] | None¶

orig_func: Callable[_P, _T] | None¶

func_annot: tilelang.language.v2.annot.FuncAnnot | None¶

out_idx_override: list[int] | None¶

class tilelang.language.v2.builder.Macro¶

Bases: Generic[_P, _T]

name: str¶

orig_func: Callable[_P, _T]¶

ir_gen: tilelang.language.v2.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.v2.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.v2.builder.get_type_hints(func)¶

tilelang.language.v2.builder.prim_func(func=None, *, generator=False)¶

Decorator to create a primitive function (PrimFunc) for TileLang IR generation. This decorator transforms a Python function into a TileLang primitive function by analyzing its type annotations and generating intermediate representation (IR) code. It supports both immediate construction (when all parameters are statically annotated) and generator mode (for dynamic construction). :param func: The function to be decorated. Can be None when using decorator with arguments. :type func: Callable[_P, _T], optional :param generator: If True, returns a generator function that creates PrimFunc instances on demand.

If False, attempts to create a PrimFunc immediately using type annotations.

Returns:

If generator=False and all parameters are statically annotated: returns a PrimFunc instance
If generator=True: returns a callable that generates PrimFunc instances when invoked
If used without parentheses: returns the decorator implementation function

Return type:

PrimFunc[_P, _T] | Callable[_P, PrimFunc[_P, _T]]

Parameters:

func (Callable[_P, _T])
generator (bool, default=False)

Examples

Static annotation mode (immediate construction): >>> @prim_func … def add_kernel(A: T.Buffer((128,), T.float32), … B: T.Buffer((128,), T.float32)): … for i in T.grid(128): … B[i] = A[i] + 1.0 Generator mode (dynamic construction): >>> @prim_func(generator=True) … def dynamic_kernel(A=T.Tensor((128,), T.float32)): … # function body … pass >>> kernel_instance = dynamic_kernel() With custom parameters: >>> @prim_func(generator=True) … def parameterized_kernel(size: int = 128): … # function body using size parameter … pass >>> kernel = parameterized_kernel(size=256)

See also

Builder: The IR builder class used for constructing primitive functions
mutate: Function used to generate IR from the decorated function