tilelang.language.customize¶

The language interface for tl programs.

Functions¶

atomic_add(dst, value)

Perform an atomic addition operation.

atomic_addx2(dst, value)

Perform an atomic addition operation with double-width operands.

atomic_addx4(dst, value)

Perform an atomic addition operation with double-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])

Views the input buffer with optionally modified shape and dtype.

Module Contents¶

tilelang.language.customize.atomic_add(dst, value)¶

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

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