tilelang.language.kernel module#

The language interface for tl programs.

class tilelang.language.kernel.FrameStack#

Bases: object

A simple stack-like wrapper around a deque that provides push, pop, and top methods for convenience.

pop()#: Pops and returns the top of the stack, or returns None if the stack is empty.

push(item)#: Pushes an item onto the top of the stack.

top()#: Returns the item on the top of the stack without removing it, or None if the stack is empty.

tilelang.language.kernel.Kernel(*blocks: List[PrimExpr], threads: Optional[Union[int, List[int], Tuple]] = None, is_cpu: bool = False, prelude: Optional[str] = None)#

Tools to quickly construct a GPU kernel launch frame.

Parameters:

blocks (List[int]) – A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
threads (int) – A integer representing blockDim.x Or a list of integers representing blockDim.(x|y|z) if the value is -1, we skip the threadIdx.x binding.
is_cpu (bool) – Whether the kernel is running on CPU. Thus we will not bind threadIdx.x, threadIdx.y, threadIdx.z. and blockIdx.x, blockIdx.y, blockIdx.z.
prelude (str) – The import c code of the kernel, will be injected before the generated kernel code.

Returns:

res – The result LaunchThreadFrame.

Return type:

Tuple[frame.LaunchThreadFrame]

class tilelang.language.kernel.KernelLaunchFrame#

Bases: TIRFrame

KernelLaunchFrame is a custom TIRFrame that manages block/thread indices and handles the entry and exit of the kernel launch scope.

classmethod Current() → Optional[KernelLaunchFrame]#: Returns the topmost (current) KernelLaunchFrame from the stack if it exists, or None if the stack is empty.

property blocks: List[Var]#: Returns the block indices from the topmost frame.

get_block_binding(dim: int = 0) → Var#: Returns the block binding for the given dimension. dim=0 corresponds to blockIdx.x, dim=1 to blockIdx.y, and dim=2 to blockIdx.z.

get_block_bindings() → List[Var]#: Returns all three block bindings.

get_block_extent(dim: int) → int#: Returns the block extent for the given dimension. dim=0 corresponds to blockIdx.x, dim=1 to blockIdx.y, and dim=2 to blockIdx.z.

get_num_threads() → int#: Returns the thread indices from the topmost frame.

get_thread_binding(dim: int = 0) → Var#: Returns the thread binding for the given dimension. dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.

get_thread_bindings() → List[Var]#: Returns the thread binding for the given dimension. dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.

get_thread_extent(dim: int) → int#: Returns the thread extent for the given dimension. dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.

property num_threads: int#: Returns the total number of threads.

property threads: List[Var]#: Returns the thread indices from the topmost frame.

tilelang.language.kernel.get_block_binding(dim: int = 0) → Var#: Returns the block binding for the given dimension.

tilelang.language.kernel.get_block_bindings() → List[Var]#: Returns all three block bindings.

tilelang.language.kernel.get_thread_binding(dim: int = 0) → Var#: Returns the thread binding for the given dimension.

tilelang.language.kernel.get_thread_bindings() → List[Var]#: Returns all three thread bindings.