tvm_tirx.backend.cuda.op.warp_reduce

Contents

tvm_tirx.backend.cuda.op.warp_reduce#

tvm_tirx.backend.cuda.op.warp_reduce(value, op, width=32, *, ty=None, loc=UNKNOWN_LOC)#

Warp-level butterfly shuffle-XOR reduction.

Reduces value across width adjacent lanes using the specified operation. Codegen emits log2(width) steps of __shfl_xor_sync(0xFFFFFFFF, val, mask) with descending XOR masks.

Parameters:

  • value (Expr) – The per-thread scalar value to reduce.

  • op (str) – Reduction operation: "sum", "max", or "min".

  • width (int) – Number of lanes participating in each reduction group. Must be a power of two in [2, 32]. Defaults to 32 (full warp).

  • loc (Location)

Returns:

call (Expr) – The reduced value (same dtype as value).

Contents