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`
`Builder`
`PrimFunc`	Abstract base class for generic types.
`Macro`	Abstract base class for generic types.

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¶

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¶

Bases: tilelang.language.v2.ast.BaseBuilder

frames: list[AnyFrame] = []¶

ir_builder¶

name_inside_frame: dict[str, AnyFrame]¶

arg_annotations¶

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

bind_immutable(name, value)¶

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

ifexp(cond, then, otherwise)¶

ret(value)¶

ctx_with(ctx)¶

assert_expr(cond, msg)¶

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

source: str | None¶

orig_func: Callable[_P, _T] | None¶

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

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