tilelang.language.v2.builder ============================ .. py:module:: tilelang.language.v2.builder Attributes ---------- .. autoapisummary:: tilelang.language.v2.builder.logger tilelang.language.v2.builder.thread_local_storage 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 Classes ------- .. autoapisummary:: tilelang.language.v2.builder.Frame tilelang.language.v2.builder.MacroFrame tilelang.language.v2.builder.ExitedMacroFrame tilelang.language.v2.builder.BoolOpFrame tilelang.language.v2.builder.ConstIfFrame tilelang.language.v2.builder.BlockFrame tilelang.language.v2.builder.ContinueFrame tilelang.language.v2.builder.BreakFrame tilelang.language.v2.builder.SerialForWithStep tilelang.language.v2.builder.Builder tilelang.language.v2.builder.PrimFunc tilelang.language.v2.builder.Macro Functions --------- .. autoapisummary:: tilelang.language.v2.builder.unwrap_expr tilelang.language.v2.builder.unwrap_cond tilelang.language.v2.builder.is_var tilelang.language.v2.builder.macro tilelang.language.v2.builder.get_type_hints tilelang.language.v2.builder.prim_func Module Contents --------------- .. py:data:: logger .. py:function:: unwrap_expr(expr) unwrap expr and convert it into PrimExpr like .. py:function:: unwrap_cond(expr) unwrap expr and convert to bool condition .. py:data:: thread_local_storage .. py:class:: Frame Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:method:: __enter__() .. py:method:: __exit__(exc_type, exc_value, traceback) .. py:class:: MacroFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: ExitedMacroFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: BoolOpFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: ConstIfFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: BlockFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: ContinueFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: BreakFrame Bases: :py:obj:`Frame` Frame are virtual context managers used in frontend only They do not have any runtime representation in the generated TIR. .. py:class:: SerialForWithStep .. py:attribute:: start :type: tvm.tir.expr.PrimExpr .. py:attribute:: stop :type: tvm.tir.expr.PrimExpr .. py:attribute:: step :type: tvm.tir.expr.PrimExpr .. py:attribute:: annotations :type: dict[str, Any] | None :value: None .. py:data:: ContinueOrBreak .. py:data:: AnyFrame .. py:data:: TIR_CONTROL_FRAME .. py:data:: TIR_VAR_SCOPE_FRAME .. py:function:: is_var(v) .. py:class:: Builder Bases: :py:obj:`tilelang.language.v2.ast.BaseBuilder` .. py:attribute:: frames :type: list[AnyFrame] :value: [] .. py:attribute:: ir_builder .. py:attribute:: name_inside_frame :type: dict[str, AnyFrame] .. py:attribute:: arg_annotations .. py:method:: current() :classmethod: .. py:method:: prim_func(name) .. py:method:: macro(name=None, annotations=None) .. py:method:: get() .. py:method:: find_frame_idx(frame, start=0) .. py:method:: enter_frame(frame) .. py:method:: check_continue_break() .. py:method:: with_frame(frame) .. py:method:: ctx_if(cond) .. py:method:: ctx_then(val) .. py:method:: ctx_else(val) .. py:method:: eval(val) .. py:method:: ctx_for(it) .. py:method:: ctx_continue() .. py:method:: ctx_break() .. py:method:: ctx_while(cond) .. py:method:: bind(name, value, annot=BaseBuilder.empty) .. py:method:: unwrap_value(value) .. py:method:: bind_immutable(name, value) .. py:method:: assign_slice(lval, sl, value, annot=BaseBuilder.empty) .. py:method:: aug_assign(op, target, aug_value) .. py:method:: aug_assign_slice(op, target, sl, aug_value) .. py:method:: boolop(op, left, right) .. py:method:: ifexp(cond, then, otherwise) .. py:method:: ret(value) .. py:method:: ctx_with(ctx) .. py:method:: assert_expr(cond, msg) .. py:method:: rval(name, value) .. py:method:: macro_arg(name, value) .. py:method:: prim_func_arg(name, value) .. py:method:: arg(name, value) .. py:method:: override(name) .. py:class:: PrimFunc Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_T`\ ], :py:obj:`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 .. py:attribute:: params :type: list[tvm.tir.Var | tvm.tir.Buffer] .. py:attribute:: body :type: tvm.tir.Stmt .. py:attribute:: ret_type :type: tvm.ir.Type .. py:attribute:: buffer_map :type: tvm_ffi.container.Map[tvm.tir.Var, tvm.tir.Buffer] .. py:attribute:: attrs :type: tvm.Attrs | None .. py:attribute:: span :type: tvm.ir.base.Span | None .. py:attribute:: ir_gen :type: tilelang.language.v2.ast.IRGenerator[_P, _T] | None .. py:attribute:: source :type: str | None .. py:attribute:: orig_func :type: Callable[_P, _T] | None .. py:class:: Macro Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_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 .. py:attribute:: name :type: str .. py:attribute:: orig_func :type: Callable[_P, _T] .. py:attribute:: ir_gen :type: tilelang.language.v2.ast.IRGenerator[_P, _T] .. py:attribute:: annotations :type: dict[str, Any] .. py:property:: source :type: str .. py:method:: __call__(*args, **kwargs) .. py:function:: 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). :type func: Callable[_P, _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]) .. seealso:: :obj:`Macro` The class that wraps macro functions :obj:`mutate` The function that transforms Python code into IR generators .. py:function:: get_type_hints(func) .. py:function:: 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. :type generator: bool, default=False :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 :rtype: PrimFunc[_P, _T] | Callable[_P, PrimFunc[_P, _T]] .. rubric:: 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) .. seealso:: :obj:`Builder` The IR builder class used for constructing primitive functions :obj:`mutate` Function used to generate IR from the decorated function