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