tilelang.language.reduce module#

The language interface for tl programs.

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

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

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, 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, 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)#

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

Returns:

Handle to the reduction operation

Return type:

tir.Call