tilelang.language.reduce_op¶

Reduce operations exposed on the TileLang language surface.

Attributes¶

ReduceKind

Functions¶

reduce(buffer, out, reduce_type, dim, clear[, batch, ...])

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

reduce_max(buffer, out[, dim, clear, batch, annotations])

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

reduce_min(buffer, out[, dim, clear, batch, annotations])

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

reduce_sum(buffer, out[, dim, clear, batch, annotations])

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

reduce_abssum(buffer, out[, dim, batch, annotations])

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

reduce_absmax(buffer, out[, dim, clear, batch, ...])

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

reduce_bitand(buffer, out[, dim, clear, batch, ...])

Perform reduce bitwise-and on input buffer, store the result to output buffer.

reduce_bitor(buffer, out[, dim, clear, batch, annotations])

Perform reduce bitwise-or on input buffer, store the result to output buffer.

reduce_bitxor(buffer, out[, dim, clear, batch, ...])

Perform reduce bitwise-xor on input buffer, store the result to output buffer.

reducer_init(reducer[, init])

Open a reducer epoch, optionally with a logical starting value.

reducer_update(target, value)

Contribute value to one logical output of a reducer.

finalize_reducer(reducer[, dst, batch, annotations])

Close a reducer epoch.

warp_reduce_sum(value)

Perform warp reduction sum on a register value.

warp_reduce_max(value)

Perform warp reduction max on a register value.

warp_reduce_min(value)

Perform warp reduction min on a register value.

warp_reduce_bitand(value)

Perform warp reduction bitwise-and on a register value.

warp_reduce_bitor(value)

Perform warp reduction bitwise-or on a register value.

Module Contents¶

tilelang.language.reduce_op.ReduceKind¶
tilelang.language.reduce_op.reduce(buffer, out, reduce_type, dim, clear, batch=1, annotations=None)¶

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

Parameters:

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

  • out (tirx.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

  • batch (int) – Number of output elements per batched AllReduce call (default 1 = scalar, current behaviour). When batch > 1 the compiler emits ceil(N/batch) batched AllReduce calls each sharing a single pair of barriers, reducing total barrier count by batch×. batch must evenly divide the per-thread output element count N.

  • annotations (dict, optional) – Additional lowering controls. The CUDA dialect exposes nan_propagate on reduce_max/min/absmax as a typed keyword (lowering to __hmax_nan/__hmin_nan); it rides here as the {"nan_propagate": True} annotation. On CUDA SM100+, FP32 sum/abssum reductions accept {"enable_fadd2": False} to keep the reducer scalar. Packed FP32x2 reduction remains enabled by default, and can be disabled globally with the tl.enable_fp32x2_reduction pass config.

Return type:

None

tilelang.language.reduce_op.reduce_max(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

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.

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • annotations (dict | None)

Returns:

handle

Return type:

PrimExpr

tilelang.language.reduce_op.reduce_min(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

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

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.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.

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • annotations (dict | None)

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reduce_sum(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

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

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.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.

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • annotations (dict, optional) – On CUDA SM100+, set {"enable_fadd2": False} to disable packed FP32x2 accumulation for this reduction. It is enabled by default, unless the tl.enable_fp32x2_reduction pass config is False.

Return type:

None

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:

tirx.Call

Parameters:

  • buffer (tvm.tirx.Buffer)

  • out (tvm.tirx.Buffer)

  • dim (int)

  • clear (bool)

  • batch (int)

  • annotations (dict | None)

tilelang.language.reduce_op.reduce_abssum(buffer, out, dim=-1, batch=1, annotations=None)¶

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

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.Buffer) – The output buffer

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

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • annotations (dict, optional) – On CUDA SM100+, set {"enable_fadd2": False} to disable packed FP32x2 accumulation for this reduction. It is enabled by default, unless the tl.enable_fp32x2_reduction pass config is False.

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reduce_absmax(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

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

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.Buffer) – The output buffer

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

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • clear (bool)

  • annotations (dict | None)

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reduce_bitand(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

Perform reduce bitwise-and on input buffer, store the result to output buffer.

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.Buffer) – The output buffer

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

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • clear (bool)

  • annotations (dict | None)

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reduce_bitor(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

Perform reduce bitwise-or on input buffer, store the result to output buffer.

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.Buffer) – The output buffer

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

  • batch (int) – Number of output elements per batched AllReduce call (default 1).

  • clear (bool)

  • annotations (dict | None)

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reduce_bitxor(buffer, out, dim=-1, clear=True, batch=1, annotations=None)¶

Perform reduce bitwise-xor on input buffer, store the result to output buffer.

Parameters:

  • buffer (tirx.Buffer) – The input buffer

  • out (tirx.Buffer) – The output buffer

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

  • clear (bool)

  • batch (int)

  • annotations (dict | None)

Returns:

Handle to the reduction operation

Return type:

tirx.Call

tilelang.language.reduce_op.reducer_init(reducer, init=None)¶

Open a reducer epoch, optionally with a logical starting value.

Must appear exactly once per T.alloc_reducer allocation, before any T.reducer_update. The whole epoch (init, updates, finalize) may sit inside thread-uniform serial loops or conditionals — the epoch then reopens once per dynamic execution — but init and finalize must share the same enclosing loop/branch scope. Without init, the reduction starts from the combine identity (sum -> 0, max -> dtype lowest, min -> dtype highest, bitand -> all ones, bitor/bitxor -> 0).

init is a LOGICAL starting value: the result is as if one extra contribution init were combined into every logical output, exactly once. It is not a physical fill — physical partials always start from the identity, and the compiler captures init at the init site and combines it once per logical output at finalize time, so physical replication can never multiply it and later writes to buffers the expression reads cannot change the epoch’s starting value.

Parameters:

  • reducer (tirx.Buffer) – Handle returned by T.alloc_reducer.

  • init (PrimExpr | int | float | None) – Optional logical starting value; converted to the reducer’s dtype when given as a Python number.

Returns:

Handle to the reducer_init intrinsic call.

Return type:

tirx.Call

tilelang.language.reduce_op.reducer_update(target, value)¶

Contribute value to one logical output of a reducer.

target must be written as acc[indices] directly in the first argument position; it is an update-target descriptor, not a read of the reducer’s current value. Each dynamic logical iteration of the enclosing T.Parallel loop contributes exactly once, regardless of how the loop is physically replicated over threads.

Parameters:

  • target (tirx.BufferLoad) – acc[indices] selecting the logical output.

  • value – Contribution expression (cast to the reducer dtype if needed).

Returns:

Handle to the reducer_update intrinsic call.

Return type:

tirx.Call

tilelang.language.reduce_op.finalize_reducer(reducer, dst=None, batch=1, annotations=None)¶

Close a reducer epoch.

v2 form (dst given): complete the cross-participant communication the chosen physical plan requires, combine the optional T.reducer_init starting value exactly once per logical output, and write the logical result into the independent destination fragment dst. After this call the reducer handle is dead; read results from dst.

Legacy v1 form (dst omitted): in-place finalize of a legacy alloc_reducer(replication=...) fragment reducer. Deprecated.

Parameters:

  • reducer (tirx.Buffer) – Reducer handle.

  • dst (tirx.Buffer | None) – Destination fragment (v2). Same logical shape and dtype as the reducer.

  • batch (int) – Batched AllReduce width: the collective covers batch output elements per call, sharing one pair of barriers.

  • annotations (dict | None)

Returns:

Handle to the finalize intrinsic call.

Return type:

tirx.Call

tilelang.language.reduce_op.warp_reduce_sum(value)¶

Perform warp reduction sum on a register value.

This function reduces a value across all threads in a warp using shuffle operations. Each thread provides a register value, and after the reduction, all threads will have the sum of all values across the warp.

Parameters:

value (tirx.PrimExpr) – The input register value to reduce

Returns:

The reduced sum value (same on all threads in the warp)

Return type:

tirx.PrimExpr

tilelang.language.reduce_op.warp_reduce_max(value)¶

Perform warp reduction max on a register value.

This function reduces a value across all threads in a warp using shuffle operations. Each thread provides a register value, and after the reduction, all threads will have the max of all values across the warp.

Parameters:

value (tirx.PrimExpr) – The input register value to reduce

Returns:

The reduced max value (same on all threads in the warp)

Return type:

tirx.PrimExpr

tilelang.language.reduce_op.warp_reduce_min(value)¶

Perform warp reduction min on a register value.

This function reduces a value across all threads in a warp using shuffle operations. Each thread provides a register value, and after the reduction, all threads will have the min of all values across the warp.

Parameters:

value (tirx.PrimExpr) – The input register value to reduce

Returns:

The reduced min value (same on all threads in the warp)

Return type:

tirx.PrimExpr

tilelang.language.reduce_op.warp_reduce_bitand(value)¶

Perform warp reduction bitwise-and on a register value.

This function reduces a value across all threads in a warp using shuffle operations. Each thread provides a register value, and after the reduction, all threads will have the bitwise-and of all values across the warp.

Parameters:

value (tirx.PrimExpr) – The input register value to reduce

Returns:

The reduced bitwise-and value (same on all threads in the warp)

Return type:

tirx.PrimExpr

tilelang.language.reduce_op.warp_reduce_bitor(value)¶

Perform warp reduction bitwise-or on a register value.

This function reduces a value across all threads in a warp using shuffle operations. Each thread provides a register value, and after the reduction, all threads will have the bitwise-or of all values across the warp.

Parameters:

value (tirx.PrimExpr) – The input register value to reduce

Returns:

The reduced bitwise-or value (same on all threads in the warp)

Return type:

tirx.PrimExpr