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

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

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]

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)¶

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.tirx.Buffer)

  • shape (tilelang._typing.ShapeType | None)

  • dtype (tilelang._typing.DType | 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