For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
WarpReducer
struct WarpReducer[WARPS_PER_BLOCK: Int = Int(1)]
Reduces a scalar across all lanes in a warp. Broadcasts the result to every lane.
Only generic uses WARPS_PER_BLOCK; the scalar paths use
hardware warp ops and are warp-count-oblivious. Default 1 keeps
the type usable in single-warp-block kernels.
Parameters
- WARPS_PER_BLOCK (
Int): Number of warps in the launching block.
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
Reducer,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__() -> Self
Default-initializes a WarpReducer.
sum
def sum[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]
Returns the sum of val across the warp.
Parameters:
- dtype (
DType): The scalar dtype.
Args:
- val (
Scalar[dtype]): The per-lane value.
Returns:
Scalar[dtype]: The warp-wide sum, broadcast to every lane.
max
def max[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]
Returns the maximum of val across the warp.
Parameters:
- dtype (
DType): The scalar dtype.
Args:
- val (
Scalar[dtype]): The per-lane value.
Returns:
Scalar[dtype]: The warp-wide maximum, broadcast to every lane.
min
def min[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]
Returns the minimum of val across the warp.
Parameters:
- dtype (
DType): The scalar dtype.
Args:
- val (
Scalar[dtype]): The per-lane value.
Returns:
Scalar[dtype]: The warp-wide minimum, broadcast to every lane.
generic
def generic[S: ReduceOp](self, mut state: S)
Warp-wide all-reduce over state.join; on return every lane holds the combined value.
Dispatches on state size (a perf choice — either path is
byte-safe): small states (<= _WARP_SHUFFLE_MAX_WORDS uint32
words) use a register-only shuffle_xor butterfly (no shmem);
larger states use a per-warp shmem tree over the warp's
[warp_id*WARP_SIZE, ...) slice, so concurrent warps don't
clobber each other.
The tree's syncwarp()s are load-bearing, not defensive: each
step reads the slot its partner lane wrote in the previous one,
and lanes diverge on lid < stride, so without them the
compiler is free to reorder the shared accesses (and Volta+
lanes need not reconverge). barrier() -- what BlockReducer's
tree uses -- is not an option here: a warp whose row_idx
exceeds num_rows skips the body entirely, so a block-wide
barrier inside it would deadlock.
Parameters:
- S (
ReduceOp): The monoid type being combined.
Args:
- state (
S): The per-lane state; on return, holds the warp-wide combined value on every lane.