tilelang.language.allocate¶

Memory allocation utilities for Tile-AI programs.

This module provides a set of functions for allocating different types of memory buffers in Tile-AI programs. It wraps TVM’s buffer allocation functionality with convenient interfaces for different memory scopes.

Available allocation functions:

alloc_shared: Allocates shared memory buffers for inter-thread communication
alloc_local: Allocates local memory buffers for thread-private storage
alloc_fragment: Allocates fragment memory buffers for specialized operations
alloc_var: Allocates single-element variable buffers
alloc_global: Allocates global memory buffers as workspace

Each function takes shape and dtype parameters and returns a TVM buffer object with the appropriate memory scope.

Attributes¶

`ReducerOp`
`DescKind`

Functions¶

`alloc_shared`(shape, dtype[, scope])	Allocate a shared memory buffer for inter-thread communication.
`alloc_local`(shape, dtype[, scope])	Allocate a local memory buffer for thread-private storage.
`alloc_fragment`(shape, dtype[, scope])	Allocate a fragment memory buffer for specialized operations.
`alloc_var`(…)	Allocate a single-element variable buffer.
`alloc_global`(shape, dtype[, scope])	Allocate a global memory buffer as a global workspace.
`alloc_barrier`(arrive_count)	Allocate a barrier buffer.
`alloc_cluster_barrier`(arrive_count)	Allocate a cluster barrier buffer.
`alloc_tmem`(shape, dtype)	Allocate a Tensor Memory (TMEM) buffer for use with 5th generation Tensor Core operations (e.g., TCGEN5.MMA).
`alloc_reducer`(shape, dtype[, op, replication])	Allocate a reducer: a first-class deferred reduction epoch handle.
`alloc_descriptor`([kind, dtype])	Allocate a descriptor buffer for WGMMA and TCGEN5.MMA.
`alloc_wgmma_desc`([dtype])
`alloc_tcgen05_smem_desc`([dtype])
`alloc_tcgen05_instruction_desc`([dtype])
`alloc_tcgen05_instr_desc`([dtype])
`empty`(shape[, dtype])	Declare the output tensor used in eager-style JIT.

Module Contents¶

tilelang.language.allocate.alloc_shared(shape, dtype, scope='shared.dyn')¶

Allocate a shared memory buffer for inter-thread communication.

Parameters:

shape (tuple) – The shape of the buffer to allocate
dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)
scope (str, optional) – The memory scope. Defaults to “shared.dyn”

Returns:

A TVM buffer object allocated in shared memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_local(shape, dtype, scope='local')¶

Allocate a local memory buffer for thread-private storage.

Parameters:

shape (tuple) – The shape of the buffer to allocate
dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)
scope (str, optional) – The memory scope. Defaults to “local”

Returns:

A TVM buffer object allocated in local memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_fragment(shape, dtype, scope='local.fragment')¶

Allocate a fragment memory buffer for specialized operations.

Parameters:

shape (tuple) – The shape of the buffer to allocate
dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)
scope (str, optional) – The memory scope. Defaults to “local.fragment”

Returns:

A TVM buffer object allocated in fragment memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_var(dtype: tilelang._typing.DType, init: tvm.tirx.PrimExpr | int | float, scope: str = 'local.var') → tvm.tirx.buffer.Buffer¶

tilelang.language.allocate.alloc_var(dtype: tilelang._typing.DType, scope: str = 'local.var', *, init: tvm.tirx.PrimExpr | int | float | None = None) → tvm.tirx.buffer.Buffer

Allocate a single-element variable buffer.

Parameters:

dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)
*args – Optional positional arguments. A single positional string is treated as the scope for backward compatibility. A single non-string positional argument (or keyword init) specifies the initializer. When two positional arguments are provided, they are interpreted as (init, scope).
scope (str, optional) – The memory scope. Defaults to “local.var”. Use as keyword argument for clarity when also providing an initializer.
init (PrimExpr, optional) – The optional initializer value. When provided, the generated code will initialize the variable with this value instead of defaulting to zero.

Examples

a = T.alloc_var(‘int32’, 1) # var with init 1 a = T.alloc_var(‘int32’, ‘local.var’) # var with local.var scope a = T.alloc_var(‘int32’, 1, ‘local.var’) # var with init 1 and local.var scope a = T.alloc_var(‘int32’, ‘local.var’, init=1) # var with init 1 and local.var scope a = T.alloc_var(‘int32’, init=1) # var with init 1 and local.var scope

Returns:: A TVM buffer object allocated as a single-element variable
Return type:: T.Buffer

tilelang.language.allocate.alloc_global(shape, dtype, scope='global')¶

Allocate a global memory buffer as a global workspace.

NOTE(chaofan): Memory allocated in this way doesn’t go through torch allocator. Instead, it’s allocated directly by the corresponding backend APIs, like cudaMalloc. We recommend allocating workspace in Torch side and pass it to the kernel via arguments, which is managed under the hood by the framework. This API is mainly for testing purposes and some specific purposes.

NOTE(chaofan): This API may not be available in all backends (e.g. CuteDSL).

Parameters:

shape (tuple) – The shape of the buffer to allocate
dtype (str) – The data type of the buffer (e.g., ‘float32’, ‘int32’)
scope (str, optional) – The memory scope. Defaults to “global”

Returns:

A TVM buffer object allocated in global memory

Return type:

T.Buffer

tilelang.language.allocate.alloc_barrier(arrive_count)¶

Allocate a barrier buffer.

Parameters:: arrive_count (int | list[int]) – The number of threads that need to arrive at each barrier. Every count must be at least 1: an mbarrier arrive count of 0 has no defined meaning, and a negative count would be reinterpreted as a garbage unsigned value at init.
Returns:: A TVM buffer object allocated as a barrier
Return type:: T.Buffer

Examples

>>> mbar = alloc_barrier(128)  # allocate a barrier with arrive count 128
>>> mbars = alloc_barrier([128] * n)  # allocate n barriers with the same arrive count 128

tilelang.language.allocate.alloc_cluster_barrier(arrive_count)¶

Allocate a cluster barrier buffer.

Parameters:: arrive_count (int | list[int]) – The number of threads that need to arrive at each barrier. Every count must be at least 1, as for alloc_barrier.
Returns:: A TVM buffer object allocated as a cluster barrier
Return type:: T.Buffer

tilelang.language.allocate.alloc_tmem(shape, dtype)¶

Allocate a Tensor Memory (TMEM) buffer for use with 5th generation Tensor Core operations (e.g., TCGEN5.MMA).

TMEM is a dedicated on-chip memory introduced in Blackwell GPUs, designed to reduce register pressure and enable asynchronous, single-threaded MMA operations. It is organized as a 2D array of 512 columns by 128 rows (lanes), with each cell being 32 bits. Allocation is performed in units of columns, and every lane of a column is allocated together.

Key properties and requirements:

The number of columns allocated must be a power of 2 and at least 32.
TMEM allocations are dynamic. TileLang deallocates them automatically at the end of the allocation block unless you call T.deallocate_tmem to take manual control of the lifetime.
Both allocation and deallocation must be performed by the same warp.
The base address of the TMEM allocation is stored in shared memory and used as the offset for TCGEN5.MMA accumulator tensors.
Only TCGEN5.MMA and specific TMEM load/store instructions can access TMEM; all pre-processing must occur before data is loaded into TMEM, and all post-processing after data is retrieved.
The number of columns allocated should not increase between any two allocations in the execution order within the CTA.

Parameters:

shape (ShapeType) – Logical buffer shape. The last two modes are the matrix modes; any leading modes are batch dimensions that repeat the buffer along TMEM columns.
dtype (DType) – Element data type.

Returns:

A TVM buffer object allocated in TMEM scope, suitable for use as an accumulator or operand in TCGEN5.MMA operations.

Return type:

T.Buffer

Note

TMEM is only available on supported architectures (e.g., Blackwell and later).
The buffer returned should be used according to TMEM access restrictions. Use T.deallocate_tmem only when you need an earlier, explicit release.

tilelang.language.allocate.ReducerOp¶

tilelang.language.allocate.alloc_reducer(shape, dtype, op='sum', replication=None)¶

Allocate a reducer: a first-class deferred reduction epoch handle.

The reducer lives in the virtual local.reducer scope and may only be accessed through the epoch operations:

acc = T.alloc_reducer(shape, dtype, op="sum")
T.reducer_init(acc)          # or T.reducer_init(acc, init_value)
for ...:
    T.reducer_update(acc[indices], contribution)
dst = T.alloc_fragment(shape, dtype)
T.finalize_reducer(acc, dst)

Ordinary reads/writes, T.clear/T.fill, aliasing, and in-place finalize are rejected at compile time. Physical storage and the cross-thread communication plan are chosen by the compiler; the physical layout can never change how many times a logical contribution is combined.

Parameters:

shape (tuple) – Logical shape of the reduction result.
dtype (str) – Element data type (e.g., ‘float32’, ‘int32’).
op (str) – Combine op: “sum”, “max”, “min”, “bitand”, “bitor” or “bitxor” (the bitwise ops require an integer dtype).
replication (str | None) – Deprecated legacy (v1) knob. Passing “all” or “none” selects the legacy fragment-based reducer for backward compatibility; it will be removed together with the v1 lowering.

Returns:

The reducer handle.

Return type:

T.Buffer

tilelang.language.allocate.DescKind¶

tilelang.language.allocate.alloc_descriptor(kind='wgmma', dtype=_dtypes.uint64)¶

Allocate a descriptor buffer for WGMMA and TCGEN5.MMA.

Parameters:

kind (DescKind) – The descriptor kind, one of “wgmma”, “tcgen05” (“utcmma” as alias).
dtype (tilelang._typing.DType)

Returns:

A TVM buffer object allocated as a descriptor

Return type:

T.Buffer

tilelang.language.allocate.alloc_wgmma_desc(dtype=_dtypes.uint64)¶

Parameters:: dtype (tilelang._typing.DType)
Return type:: tvm.tirx.buffer.Buffer

tilelang.language.allocate.alloc_tcgen05_smem_desc(dtype=_dtypes.uint64)¶

Parameters:: dtype (tilelang._typing.DType)
Return type:: tvm.tirx.buffer.Buffer

tilelang.language.allocate.alloc_tcgen05_instruction_desc(dtype=_dtypes.uint32)¶

Parameters:: dtype (tilelang._typing.DType)
Return type:: tvm.tirx.buffer.Buffer

tilelang.language.allocate.alloc_tcgen05_instr_desc(dtype=_dtypes.uint32)¶

Parameters:: dtype (tilelang._typing.DType)
Return type:: tvm.tirx.buffer.Buffer

tilelang.language.allocate.empty(shape, dtype=_dtypes.float32)¶

Declare the output tensor used in eager-style JIT.

Tensors allocated in this way should be returned as the output of the function.

Parameters:

shape (tuple) – The shape of the tensor to allocate
dtype (str) – The data type of the tensor (e.g., ‘float32’, ‘int32’)

Returns:

The declared OutTensor object.

Return type:

Tensor