tilelang.language.customize¶

Some customized operations frequently used in tensor programming, exposed on the TileLang language surface.

Functions¶

dp4a(A, B, C)

Perform a four-element signed int8 dot product accumulated into int32.

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, strides])

Return a Tensor view with an optional new shape, dtype, and explicit strides.

loop_break()

Break out of the current loop.

Module Contents¶

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

Perform a four-element signed int8 dot product accumulated into int32.

Parameters:

  • A (tilelang._typing.BufferLikeType) – First int8 input buffer.

  • B (tilelang._typing.BufferLikeType) – Second int8 input buffer.

  • C (tilelang._typing.BufferLikeType) – Int32 accumulator buffer.

Returns:

Handle to the DP4A operation.

Raises:

ValueError – If A or B is not int8, or C is not int32.

Return type:

tvm.tirx.PrimExpr

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

Clamps the input value dst between [min_val, max_val]

Equivalent to T.min(T.max(dst, min_val), max_val). NaN behavior follows the target’s min/max operations without additional checks.

Parameters:

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

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

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

Returns:

Value clamped to the specified range

Return type:

tvm.tirx.PrimExpr

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

Reshapes the input buffer to the specified shape.

Parameters:

  • src (Buffer) – Input buffer to be reshaped

  • shape (ShapeType) – 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, strides=None)¶

Return a Tensor view with an optional new shape, dtype, and explicit strides.

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 its data with src without copying.

Without strides, src must be densely packed and the result uses a dense layout. A non-contiguous source requires explicit strides. Explicit strides are expressed in units of dtype elements and allow rank changes.

The source must have zero elem_offset. The caller must guarantee that explicit strides are valid for the backing storage, including dtype alignment and physical contiguity between any source dimensions merged by a rank-changing view.

Parameters:

  • src (tvm.tirx.Buffer)

  • shape (tilelang._typing.ShapeType | None)

  • dtype (tilelang._typing.DType | None)

  • strides (tilelang._typing.ShapeType | None)

Return type:

tvm.tirx.Buffer

tilelang.language.customize.loop_break()¶

Break out of the current loop.

Returns:

A call to the tl.loop_break intrinsic.

Return type:

tir.Call