tilelang.language.tir.entry module#

tilelang.language.tir.entry.macro(*args, hygienic: bool = True) → Callable#

Decorator for macro definitions.

Parameters:

hygienic –

Specifies whether the macro is hygienic or not. A macro is hygienic if all symbols used in the macro’s body are resolved to values from the location of the macro definition. A non-hygienic macro will have its symbols resolved to values at the time of the macro’s use.

Example: ``` import tvm from tvm.script import tir as T

x_value = 128

@T.macro(hygienic=True) def static_capture(A, B):

B[()] = A[x_value] ### x_value binds to 128

@T.macro(hygienic=False) def dynamic_capture(A, B):

B[()] = A[x_value] ### x_value will bind at the time of use

tilelang.language.tir.entry.prim_func(func: Optional[Callable] = None, private: bool = False, check_well_formed=True) → Union[PrimFunc, Callable]#

The parsing method for tir prim func, by using @prim_func as decorator.

Parameters:

  • func (Callable) – The function to be parsed as prim func. (Listed as optional to allow the decorator to be used without arguments, like @prim_func, or with an argument, @prim_func(private=True))

  • private (bool, optional) – Whether the function should be treated as private. A private function has no global symbol attribute; if the function is not private, it will have a global symbol matching the function name.

Returns:

res – The parsed tir prim func.

Return type:

Union[PrimFunc, Callable]