tvm_tirx.backend.cuda.op.cta_reduce

Contents

tvm_tirx.backend.cuda.op.cta_reduce#

tvm_tirx.backend.cuda.op.cta_reduce(value, op, num_warps, scratch, *, ty=None, loc=UNKNOWN_LOC)#

CTA-wide reduction via warp shuffle + shared memory.

Two-step reduction: (1) intra-warp shuffle reduction, (2) warp-0 collects per-warp partials from scratch, reduces, broadcasts via __syncthreads(). All CTA threads must participate.

Parameters:

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

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

  • num_warps (int) – Number of warps in the CTA. Must be a power of two in [1, 32].

  • scratch (Var) – Data pointer to shared-memory scratch space (>= num_warps elements).

  • loc (Location)

Returns:

call (Expr) – The reduced value broadcast to all threads (same dtype as value).

Contents