tilelang.language.reduce module#

The language interface for tl programs.

tilelang.language.reduce.cumsum(src: Buffer, dst: Optional[Buffer] = None, dim: int = 0, reverse: bool = False)#

Perform cumulative sum on input buffer, store the result to output buffer.

Parameters:

  • src (tir.Buffer) – The input buffer

  • dst (tir.Buffer, optional) – The output buffer. Defaults to None.

  • dim (int, optional) – The dimension to perform cumulative sum on. Defaults to 0.

  • reverse (bool, optional) – Whether to perform reverse cumulative sum. Defaults to False.

Returns:

Handle to the cumulative sum operation

Return type:

tir.Call

tilelang.language.reduce.reduce(buffer: Buffer, out: Buffer, reduce_type: str, dim: int, clear: bool)#

Perform a reduction operation on a buffer along a specified dimension.

Parameters:

  • buffer (tir.Buffer) – Input buffer to reduce

  • out (tir.Buffer) – Output buffer to store results

  • reduce_type (str) – Type of reduction (‘max’, ‘min’, ‘sum’, ‘abssum’)

  • dim (int) – Dimension along which to perform reduction

  • clear (bool) – Whether to initialize the output buffer before reduction

Returns:

Handle to the reduction operation

Return type:

tir.Call

tilelang.language.reduce.reduce_absmax(buffer: Buffer, out: Buffer, dim: int = -1, clear: bool = True)#

Perform reduce absolute max on input buffer, store the result to output buffer.

Parameters:

  • buffer (tir.Buffer) – The input buffer

  • out (tir.Buffer) – The output buffer

  • dim (int) – The dimension to perform reduce on

Returns:

Handle to the reduction operation

Return type:

tir.Call

tilelang.language.reduce.reduce_abssum(buffer: Buffer, out: Buffer, dim: int = -1)#

Perform reduce absolute sum on input buffer, store the result to output buffer.

Parameters:

  • buffer (tir.Buffer) – The input buffer

  • out (tir.Buffer) – The output buffer

  • dim (int) – The dimension to perform reduce on

Returns:

Handle to the reduction operation

Return type:

tir.Call

tilelang.language.reduce.reduce_max(buffer: Buffer, out: Buffer, dim: int = -1, clear: bool = True)#

Perform reduce max on input buffer, store the result to output buffer

Parameters:

  • buffer (Buffer) – The input buffer.

  • out (Buffer) – The output buffer.

  • dim (int) – The dimension to perform reduce on

  • clear (bool) – If set to True, the output buffer will first be initialized to -inf.

Returns:

handle

Return type:

PrimExpr

tilelang.language.reduce.reduce_min(buffer: Buffer, out: Buffer, dim: int = -1, clear: bool = True)#

Perform reduce min on input buffer, store the result to output buffer.

Parameters:

  • buffer (tir.Buffer) – The input buffer

  • out (tir.Buffer) – The output buffer

  • dim (int) – The dimension to perform reduce on

  • clear (bool, optional) – If True, output buffer will be initialized to inf. Defaults to True.

Returns:

Handle to the reduction operation

Return type:

tir.Call

tilelang.language.reduce.reduce_sum(buffer: Buffer, out: Buffer, dim: int = -1, clear: bool = True)#

Perform reduce sum on input buffer, store the result to output buffer.

Parameters:

  • buffer (tir.Buffer) – The input buffer

  • out (tir.Buffer) – The output buffer

  • dim (int) – The dimension to perform reduce on

  • clear (bool, optional) – If True, output buffer will be cleared before reduction. If False, results will be accumulated on existing values. Defaults to True.

Note: When clear=True, reduce_sum will not compute directly on the output buffer. This is because

during warp reduction, the same value would be accumulated multiple times (number of threads in the warp). Therefore, the implementation with clear=True follows these steps:

  1. create a temp buffer with same shape and dtype as out

  2. copy out to temp buffer

  3. call reduce_sum with temp buffer and out

  4. Add temp buffer to out

Returns:

Handle to the reduction operation

Return type:

tir.Call