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.OutTensor tilelang.language.v2.builder.Ref tilelang.language.v2.builder.UnrollForWithStep tilelang.language.v2.builder.Builder tilelang.language.v2.builder.PrimFuncCreater 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:class:: OutTensor .. py:attribute:: shape :type: collections.abc.Sequence[tvm.tir.expr.PrimExpr] .. py:attribute:: dtype :type: tilelang.language.v2.dtypes.dtype .. py:property:: strides .. py:class:: Ref .. py:attribute:: bufload :type: tvm.tir.expr.BufferLoad .. py:property:: buffer .. py:method:: store(value) .. py:method:: load() .. py:class:: UnrollForWithStep Bases: :py:obj:`SerialForWithStep` .. 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(func_annot = None) 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:: macro_arg_annot .. py:attribute:: func_annot :value: None .. py:attribute:: out_idx :value: [] .. py:attribute:: out_tensor_cnt :value: 0 .. py:attribute:: arg_vt .. 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) Unwrap some tilelang objects to get their inner value .. py:method:: 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. .. 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=None) .. py:method:: ifexp(cond, then, otherwise) .. py:method:: ret(value=None) .. py:method:: ctx_with(ctx) .. py:method:: assert_expr(cond, msg=None) .. 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:: PrimFuncCreater Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_T`\ ] .. py:attribute:: func_annot :type: tilelang.language.v2.annot.FuncAnnot .. py:attribute:: ir_gen :type: tilelang.language.v2.ast.IRGenerator[_P, _T] .. py:attribute:: orig_func :type: Callable[_P, _T] .. py:property:: annot :type: dict[str, tilelang.language.v2.annot.Annot] .. py:method:: __call__(*args, **kwargs) .. py:method:: __repr__() .. py:class:: PrimFunc Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_T`\ ], :py:obj:`tvm.tir.PrimFunc` .. 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:: orig_func :type: Callable[_P, _T] | None .. py:attribute:: func_annot :type: tilelang.language.v2.annot.FuncAnnot | None .. py:attribute:: out_idx_override :type: list[int] | None .. py:class:: Macro Bases: :py:obj:`Generic`\ [\ :py:obj:`_P`\ , :py:obj:`_T`\ ] .. 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:method:: __hash__() .. py:method:: __eq__(other) .. 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:: :py:obj:`Macro` The class that wraps macro functions :py: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:: :py:obj:`Builder` The IR builder class used for constructing primitive functions :py:obj:`mutate` Function used to generate IR from the decorated function