tilelang.language.kernel¶

Kernel launching language interface in TileLang.

Classes¶

`FrameStack`	A simple stack-like wrapper around a deque that provides
`KernelLaunchFrame`	KernelLaunchFrame is a custom TIRFrame that manages block/thread indices

Functions¶

`kernel_launch_factory`(func)	Mark `func` as a launch factory: a callable used as `with func(...)`
`is_kernel_launch_factory`(obj)
`launch_kernel`(blocks, *[, threads, prelude, cluster_dims])	Shared implementation behind every dialect's `T.Kernel`.
`Kernel`(*blocks)	Construct a kernel launch frame: a grid of tile programs.
`ClusterKernel`(*blocks, cluster_dims[, threads, prelude])	Construct a kernel launch frame with a CUDA thread block cluster
`CUDASourceCodeKernel`(*blocks[, threads, entry_name, ...])	Launch a kernel from CUDA source code or a CUDA source file.
`get_thread_binding`([dim])	Returns the thread binding for the given dimension.
`get_thread_bindings`()	Returns all three thread bindings.
`get_block_binding`([dim])	Returns the block binding for the given dimension.
`get_block_bindings`()	Returns all three block bindings.
`get_thread_extent`([dim])	Returns the thread extent for the given dimension.
`get_thread_extents`()	Returns all three thread extents.
`get_block_extent`([dim])	Returns the block extent for the given dimension.
`get_block_extents`()	Returns all three block extents.
`get_cluster_dims`()	Returns the cluster dimensions `[x, y, z]` of the current launch (`[1, 1, 1]` without clusters).
`get_cluster_size`()	Returns the number of programs per cluster of the current launch.
`get_cluster_id`([dim])	Returns the cluster index of the current program along `dim`
`get_cluster_ids`()	Returns the cluster index along every launched grid axis.
`get_cluster_extent`([dim])	Returns the number of clusters along `dim`.
`get_cluster_extents`()	Returns the number of clusters along all three dimensions.

Module Contents¶

class tilelang.language.kernel.FrameStack¶

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

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

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

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

size()¶: Returns the number of items in the stack.

__len__()¶: Returns the number of items in the stack.

__bool__()¶: Allows truthy checks on the stack object itself, e.g., ‘if stack: …’

class tilelang.language.kernel.KernelLaunchFrame¶

Bases: tvm.tirx.script.builder.frame.TIRFrame

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

Grid (program index) vars are bound by the frame itself. Thread vars are placeholders: they have an identity so the body can reference them, but their extent is only known once a backend materializes the launch. Thread extents are therefore available at trace time only when threads= was passed to Kernel().

__enter__()¶

Enters the KernelLaunchFrame scope and pushes this frame onto the stack. Returns one Var for a single grid dimension, or a list of Vars otherwise.

Return type:: tvm.tirx.Var | list[tvm.tirx.Var]

__exit__(ptype, value, trace)¶: Exits the KernelLaunchFrame scope and pops this frame from the stack, but only if it’s indeed the topmost frame.

classmethod Current()¶

Returns the topmost (current) KernelLaunchFrame from the stack if it exists, or None if the stack is empty.

Return type:: KernelLaunchFrame | None

get_block_extent(dim)¶

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. Grid axes that were not launched have extent 1.

Parameters:: dim (int)
Return type:: int

get_block_extents()¶

Returns the block extents for all three dimensions.

Return type:: list[int]

get_thread_extent(dim)¶

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.

Raises:: ValueError – If the kernel was launched without threads=. The extent is then chosen by the backend and is not known at trace time.
Parameters:: dim (int)
Return type:: int

get_thread_extents()¶

Returns the thread extents for all three dimensions.

Return type:: list[int]

get_thread_binding(dim=0)¶

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.

Parameters:: dim (int)
Return type:: tvm.tirx.Var

get_thread_bindings()¶

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.

Return type:: list[tvm.tirx.Var]

get_num_threads()¶

Returns the thread indices from the topmost frame.

Return type:: int

get_block_binding(dim=0)¶

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.

Parameters:: dim (int)
Return type:: tvm.tirx.Var

get_block_bindings()¶

Returns all three block bindings.

Return type:: list[tvm.tirx.Var]

get_launch_annotation(key, default=None)¶

Returns the launch annotation key recorded by T.Kernel (e.g. cluster_dims), or default when it was not given.

Parameters:: key (str)

get_cluster_dims()¶

Returns the cluster dimensions as [x, y, z]. A launch without cluster_dims has clusters of a single program, i.e. [1, 1, 1].

Return type:: list[int]

get_cluster_size()¶

Returns the number of programs per cluster (product of the cluster dimensions).

Return type:: int

get_cluster_id(dim=0)¶

Returns the index of the cluster the current program belongs to along dim, in program-space arithmetic: block_id // cluster_dims[dim].

A cluster is a cluster_dims-shaped tile of the grid, so this is the same on every target (clusterIdx on CUDA, a group of consecutive programs elsewhere) and stays consistent with threadblock swizzling, which permutes the grid at cluster granularity.

Parameters:: dim (int)
Return type:: tvm.tirx.Var | tvm.tirx.PrimExpr

get_cluster_ids()¶

Returns the cluster index along every launched grid axis.

Return type:: list[tvm.tirx.Var | tvm.tirx.PrimExpr]

get_cluster_extent(dim=0)¶

Returns the number of clusters along dim: ceil(grid_extent / cluster_dims[dim]).

Parameters:: dim (int)
Return type:: int

get_cluster_extents()¶

Returns the number of clusters along all three dimensions.

Return type:: list[int]

property blocks: list[tvm.tirx.Var]¶

Returns the block indices from the topmost frame.

Return type:: list[tvm.tirx.Var]

property threads: list[tvm.tirx.Var]¶

Returns the thread indices from the topmost frame.

Return type:: list[tvm.tirx.Var]

property num_threads: int¶

Returns the total number of threads.

Return type:: int

tilelang.language.kernel.kernel_launch_factory(func)¶: Mark func as a launch factory: a callable used as with func(...) to open a kernel launch. Every dialect’s Kernel (and ClusterKernel) carries this mark so the eager JIT rewriter can find the launch regardless of which dialect or alias the user went through.

tilelang.language.kernel.is_kernel_launch_factory(obj)¶

Return type:: bool

tilelang.language.kernel.launch_kernel(blocks, *, threads=None, prelude=None, cluster_dims=None, **annotations)¶

Shared implementation behind every dialect’s T.Kernel.

The well-known launch annotations are normalized here; any other keyword a dialect forwards is recorded verbatim on the launch block for that backend’s pipeline to consume. Dialects, not this function, decide which keywords exist: they only forward what their own Kernel signature declares.

Parameters:

blocks (tuple[int | tvm.tirx.PrimExpr, Ellipsis])
threads (int | list[int] | tuple[int, Ellipsis] | None)
prelude (str | None)
cluster_dims (int | tuple[int, int, int] | list[int] | None)
annotations (Any)

Return type:

KernelLaunchFrame

tilelang.language.kernel.Kernel(*blocks)¶

Construct a kernel launch frame: a grid of tile programs.

This is the target-neutral launch: the part every backend shares. Backend dialects offer their own T.Kernel with the launch annotations that backend understands, e.g. tilelang.cuda.language.Kernel(..., threads=128); tilelang.language is the CUDA dialect.

Parameters:: *blocks (int | PrimExpr) – Extent of the grid along each axis (1-3 dimensions). The launch yields one program index per axis (blockIdx on CUDA, the outer loop on CPU, the core index on an NPU).
Return type:: KernelLaunchFrame

Examples

with T.Kernel(T.ceildiv(N, 128)) as bx:
    # bx is the program index along x; also iterable as (bx,)
    ...

tilelang.language.kernel.ClusterKernel(*blocks, cluster_dims, threads=None, prelude=None)¶

Construct a kernel launch frame with a CUDA thread block cluster (SM90+ only).

This is the CUDA-specific variant of Kernel(): identical launch semantics and bindings, plus a cluster_dims annotation. The kernel will be launched with cudaLaunchKernelEx using cudaLaunchAttributeClusterDimension.

Parameters:

blocks (int) – A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
cluster_dims (int | tuple[int, int, int] | list[int]) – The cluster dimensions. For example, use 2 or (2, 1, 1) to create 2-CTA clusters.
threads (int) – A integer representing blockDim.x Or a list of integers representing blockDim.(x|y|z)
prelude (str) – The import c code of the kernel, will be injected before the generated kernel code.

Examples

with T.ClusterKernel(grid_x, grid_y, cluster_dims=2, threads=128) as (bx, by):
    ...

tilelang.language.kernel.CUDASourceCodeKernel(*blocks, threads=None, source_code_or_path, entry_name='main_kernel', cluster_dims=None, prelude=None)¶

Launch a kernel from CUDA source code or a CUDA source file.

The code must follows the following rules: 1. The kernel source must be a valid CUDA kernel which can be correctly compiled under TileLang’s context. 2. The kernel source must either contains only one __global__ function as an entry, or have a __global__ entry function named main_kernel.

Parameters:

source_code_or_path (str | os.PathLike[str]) – Inline CUDA source code, or a path to a CUDA source file. If the argument resolves to an existing file, the file contents are loaded. Otherwise it is treated as inline CUDA source code.
blocks (int) – A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
entry_name (str | None) – Optional name of the __global__ CUDA entry function inside the provided source. When specified, TileLang launches that external CUDA entry directly.
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.
cluster_dims (int | tuple[int, int, int] | list[int] | None) – The cluster dimensions for SM90+ cluster launch. For example, use 2 or (2, 1, 1) to create 2-CTA clusters. When specified, the kernel will be launched using cudaLaunchKernelEx with cudaLaunchAttributeClusterDimension.
prelude (str) – The import c code of the kernel, will be injected before the generated kernel code.

Return type:

None

tilelang.language.kernel.get_thread_binding(dim=0)¶

Returns the thread binding for the given dimension.

Parameters:: dim (int)
Return type:: tvm.tirx.Var

tilelang.language.kernel.get_thread_bindings()¶

Returns all three thread bindings.

Return type:: list[tvm.tirx.Var]

tilelang.language.kernel.get_block_binding(dim=0)¶

Returns the block binding for the given dimension.

Parameters:: dim (int)
Return type:: tvm.tirx.Var

tilelang.language.kernel.get_block_bindings()¶

Returns all three block bindings.

Return type:: list[tvm.tirx.Var]

tilelang.language.kernel.get_thread_extent(dim=0)¶

Returns the thread extent for the given dimension.

Parameters:: dim (int)
Return type:: int

tilelang.language.kernel.get_thread_extents()¶

Returns all three thread extents.

Return type:: list[int]

tilelang.language.kernel.get_block_extent(dim=0)¶

Returns the block extent for the given dimension.

Parameters:: dim (int)
Return type:: int

tilelang.language.kernel.get_block_extents()¶

Returns all three block extents.

Return type:: list[int]

tilelang.language.kernel.get_cluster_dims()¶

Returns the cluster dimensions [x, y, z] of the current launch ([1, 1, 1] without clusters).

Return type:: list[int]

tilelang.language.kernel.get_cluster_size()¶

Returns the number of programs per cluster of the current launch.

Return type:: int

tilelang.language.kernel.get_cluster_id(dim=0)¶

Returns the cluster index of the current program along dim (block_id // cluster_dims[dim]). See KernelLaunchFrame.get_cluster_id().

Parameters:: dim (int)
Return type:: tvm.tirx.Var | tvm.tirx.PrimExpr

tilelang.language.kernel.get_cluster_ids()¶

Returns the cluster index along every launched grid axis.

Return type:: list[tvm.tirx.Var | tvm.tirx.PrimExpr]

tilelang.language.kernel.get_cluster_extent(dim=0)¶

Returns the number of clusters along dim.

Parameters:: dim (int)
Return type:: int

tilelang.language.kernel.get_cluster_extents()¶

Returns the number of clusters along all three dimensions.

Return type:: list[int]