tilelang.language.customize

The language interface for tl programs.

Functions

region(buffer, access_type, *args)

Create a tile memory-region descriptor for a BufferLoad.

buffer_to_tile_region(buffer, access_type)

Convert a TVM buffer to a tile region descriptor.

buffer_load_to_tile_region(load, access_type, extents)

Convert a buffer load operation to a tile region descriptor.

buffer_region_to_tile_region(buffer_region, ...)

Create a tl region descriptor for the given BufferRegion.

atomic_max(dst, value[, memory_order])

Perform an atomic maximum on the value stored at dst with an optional memory-order.

atomic_min(dst, value[, memory_order])

Atomically update the value at dst to the minimum of its current value and value.

atomic_add(dst, value[, memory_order])

Atomically add value into dst, returning a handle to the operation.

atomic_addx2(dst, value)

Perform an atomic addition operation with double-width operands.

atomic_addx4(dst, value)

Perform an atomic addition operation with quad-width operands.

dp4a(A, B, C)

Perform a 4-element dot product with accumulation (DP4A).

clamp(dst, min_val, max_val)

Clamps the input value dst between [min_val, max_val]

reshape(src, shape)

Reshapes the input buffer to the specified shape.

view(src[, shape, dtype])

Return a Tensor view of the input buffer with an optional new shape and dtype.

atomic_load(src[, memory_order])

Load a value from the given buffer using the specified atomic memory ordering.

atomic_store(dst, src[, memory_order])

Perform an atomic store of src into dst with the given memory ordering.

Module Contents

tilelang.language.customize.region(buffer, access_type, *args)

Create a tile memory-region descriptor for a BufferLoad.

Maps access_type (‘r’, ‘w’, ‘rw’) to the numeric codes expected by the tl.region intrinsic (1, 2, 3 respectively) and returns a tir.Call representing the region with the provided extents.

Parameters:

  • buffer (tir.BufferLoad) – The BufferLoad that identifies the underlying buffer and indices.

  • access_type (str) – One of ‘r’, ‘w’, or ‘rw’ indicating read, write, or read-write access.

  • *args (tir.PrimExpr) – Extent expressions for each region dimension.

Returns:

A call to the tl.region intrinsic describing the memory region.

Return type:

tir.Call

Raises:

KeyError – If access_type is not one of ‘r’, ‘w’, or ‘rw’.

tilelang.language.customize.buffer_to_tile_region(buffer, access_type)

Convert a TVM buffer to a tile region descriptor.

Parameters:

  • buffer (tir.Buffer) – The buffer to convert

  • access_type (str) – Type of access - ‘r’ for read, ‘w’ for write, ‘rw’ for read-write

Returns:

A region descriptor covering the entire buffer

Return type:

tir.Call

tilelang.language.customize.buffer_load_to_tile_region(load, access_type, extents)

Convert a buffer load operation to a tile region descriptor.

Parameters:

  • load (tir.BufferLoad) – The buffer load operation

  • access_type (str) – Type of access - ‘r’ for read, ‘w’ for write, ‘rw’ for read-write

  • extents (List[tir.PrimExpr]) – List of expressions defining the region size

Returns:

A region descriptor for the loaded area

Return type:

tir.Call

tilelang.language.customize.buffer_region_to_tile_region(buffer_region, access_type, extents)

Create a tl region descriptor for the given BufferRegion.

Parameters:

  • buffer_region (tir.BufferRegion) – Source buffer region whose region items provide mins and extents.

  • access_type (str) – Access mode: “r”, “w”, or “rw”.

  • extents (List[PrimExpr]) – Requested extents; must have length <= the number of extents in buffer_region.region.

Returns:

A tile-region descriptor (tl.region) covering the buffer_region.

Return type:

tir.Call

Raises:

AssertionError – If the number of extents in buffer_region.region is smaller than len(extents).

tilelang.language.customize.atomic_max(dst, value, memory_order=None)

Perform an atomic maximum on the value stored at dst with an optional memory-order.

If memory_order is None the runtime extern “AtomicMax” is called without an explicit memory-order id; otherwise the provided memory_order string is mapped to a numeric id using the module’s memory-order map and passed to the extern.

Parameters:

  • dst (Buffer) – Destination buffer/address to apply the atomic max.

  • value (PrimExpr) – Value to compare/store atomically.

  • memory_order (str | None) – Optional memory-order name (e.g. “relaxed”, “acquire”, “seq_cst”). If provided, it is translated to the corresponding numeric memory-order id before the call.

Returns:

A handle/expression representing the issued atomic maximum operation.

Return type:

PrimExpr

tilelang.language.customize.atomic_min(dst, value, memory_order=None)

Atomically update the value at dst to the minimum of its current value and value.

If memory_order is provided, it selects the memory-order semantic used by the underlying extern call; allowed names are “relaxed”, “consume”, “acquire”, “release”, “acq_rel”, and “seq_cst” (mapped internally to integer IDs). If memory_order is None, the extern is invoked without an explicit memory-order argument.

Parameters:

  • memory_order (str | None) – Optional memory-order name controlling the atomic operation’s ordering.

  • dst (tvm.tir.Buffer)

  • value (tvm.tir.PrimExpr)

Returns:

A handle expression representing the atomic-min operation.

Return type:

PrimExpr

tilelang.language.customize.atomic_add(dst, value, memory_order=None)

Atomically add value into dst, returning a handle to the operation.

Supports scalar/addressed extern atomic add when neither argument exposes extents, or tile-region-based atomic add for Buffer/BufferRegion/BufferLoad inputs. If both arguments are plain Buffers their shapes must be structurally equal. If at least one side exposes extents, extents are aligned (missing dimensions are treated as size 1); an assertion is raised if extents cannot be deduced. The optional memory_order (one of “relaxed”,”consume”,”acquire”,”release”,”acq_rel”,”seq_cst”) is used only for the direct extern AtomicAdd path when no extents are available — otherwise the tile-region path ignores memory_order.

Returns:

A handle representing the atomic addition operation.

Return type:

PrimExpr

Parameters:

  • dst (tvm.tir.Buffer)

  • value (tvm.tir.PrimExpr)

  • memory_order (str | None)

tilelang.language.customize.atomic_addx2(dst, value)

Perform an atomic addition operation with double-width operands.

Parameters:

  • dst (Buffer) – Destination buffer where the atomic addition will be performed

  • value (PrimExpr) – Value to be atomically added (double-width)

Returns:

Handle to the double-width atomic addition operation

Return type:

PrimExpr

tilelang.language.customize.atomic_addx4(dst, value)

Perform an atomic addition operation with quad-width operands.

Parameters:

  • dst (Buffer) – Destination buffer where the atomic addition will be performed

  • value (PrimExpr) – Value to be atomically added (quad-width)

Returns:

Handle to the quad-width atomic addition operation

Return type:

PrimExpr

tilelang.language.customize.dp4a(A, B, C)

Perform a 4-element dot product with accumulation (DP4A).

Parameters:

  • A (Buffer) – First input buffer

  • B (Buffer) – Second input buffer

  • C (Buffer) – Accumulation buffer

Returns:

Handle to the DP4A operation

Return type:

PrimExpr

tilelang.language.customize.clamp(dst, min_val, max_val)

Clamps the input value dst between [min_val, max_val]

Parameters:

  • dst (tvm.tir.PrimExpr) – Input value to be clamped

  • min_val (tvm.tir.PrimExpr) – Minimum value

  • max_val (tvm.tir.PrimExpr) – Maximum value

Returns:

Value clamped to the specified range

Return type:

tvm.tir.PrimExpr

tilelang.language.customize.reshape(src, shape)

Reshapes the input buffer to the specified shape.

Parameters:

  • src (Buffer) – Input buffer to be reshaped

  • shape (List[PrimExpr]) – New shape for the buffer

Returns:

A new buffer view with the specified shape

Return type:

Buffer

tilelang.language.customize.view(src, shape=None, dtype=None)

Return a Tensor view of the input buffer with an optional new shape and dtype.

If shape is None the source buffer’s shape is used; if dtype is None the source buffer’s dtype is used. The returned buffer shares the same underlying data as src (no copy).

Parameters:

  • src (tvm.tir.Buffer)

  • shape (Union[List[tvm.tir.PrimExpr], None])

  • dtype (Union[str, None])

Return type:

tvm.tir.Buffer

tilelang.language.customize.atomic_load(src, memory_order='seq_cst')

Load a value from the given buffer using the specified atomic memory ordering.

Performs an atomic load from src and returns a PrimExpr representing the loaded value. memory_order selects the ordering and must be one of: “relaxed”, “consume”, “acquire”, “release”, “acq_rel”, or “seq_cst” (default). Raises KeyError if an unknown memory_order is provided.

Parameters:

  • src (tvm.tir.Buffer)

  • memory_order (str)

Return type:

tvm.tir.PrimExpr

tilelang.language.customize.atomic_store(dst, src, memory_order='seq_cst')

Perform an atomic store of src into dst with the given memory ordering.

Parameters:

  • dst (Buffer) – Destination buffer to store into.

  • src (PrimExpr) – Value to store.

  • memory_order (str, optional) – Memory ordering name; one of “relaxed”, “consume”, “acquire”, “release”, “acq_rel”, or “seq_cst”. Defaults to “seq_cst”. The name is mapped to an internal numeric ID used by the underlying runtime.

Returns:

A handle representing the issued atomic store operation.

Return type:

PrimExpr

Raises:

KeyError – If memory_order is not one of the supported names.