tilelang.language.customize module#

The language interface for tl programs.

tilelang.language.customize.atomic_add(dst: Buffer, value: PrimExpr) → PrimExpr#

Perform an atomic addition operation.

Parameters:

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

  • value (PrimExpr) – Value to be atomically added

Returns:

Handle to the atomic addition operation

Return type:

PrimExpr

tilelang.language.customize.atomic_addx2(dst: Buffer, value: PrimExpr) → PrimExpr#

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.clamp(dst: PrimExpr, min_val: PrimExpr, max_val: PrimExpr) → PrimExpr#

Clamps the input value dst between [min_val, max_val]

Parameters:

  • dst – Input value to be clamped

  • min_val – Minimum value

  • max_val – Maximum value

Returns:

Value clamped to the specified range

tilelang.language.customize.dp4a(A: Buffer, B: Buffer, C: Buffer) → PrimExpr#

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.reshape(src: Buffer, shape: List[PrimExpr]) → Buffer#

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: Buffer, shape: Optional[List[PrimExpr]] = None, dtype: Optional[str] = None) → Buffer#

Views the input buffer with optionally modified shape and dtype.

Parameters:

  • src (Buffer) – Input buffer to be viewed

  • shape (Union[List[PrimExpr], None], optional) – New shape for the buffer. Defaults to None.

  • dtype (Union[str, None], optional) – New dtype for the buffer. Defaults to None.

Returns:

A new buffer view with the specified shape and dtype

Return type:

Buffer