IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

BlockReducer

struct BlockReducer[BLOCK_SIZE: Int]

Reduces a scalar across BLOCK_SIZE threads in a block. Broadcasts the result to every thread.

Parameters​

  • ​BLOCK_SIZE (Int): Number of threads in the launching block.

Implemented traits​

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, Reducer, RegisterPassable, TrivialRegisterPassable

Methods​

__init__​

def __init__() -> Self

Default-initializes a BlockReducer.

sum​

def sum[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]

Returns the sum of val across the block.

Parameters:

  • ​dtype (DType): The scalar dtype.

Args:

Returns:

Scalar[dtype]: The block-wide sum, broadcast to every thread.

max​

def max[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]

Returns the maximum of val across the block.

Parameters:

  • ​dtype (DType): The scalar dtype.

Args:

Returns:

Scalar[dtype]: The block-wide maximum, broadcast to every thread.

min​

def min[dtype: DType](self, val: Scalar[dtype]) -> Scalar[dtype]

Returns the minimum of val across the block.

Parameters:

  • ​dtype (DType): The scalar dtype.

Args:

Returns:

Scalar[dtype]: The block-wide minimum, broadcast to every thread.

generic​

def generic[S: ReduceOp](self, mut state: S)

Block-wide all-reduce over state.join; on return every thread holds the combined value.

Dispatches on state size (a perf choice β€” either path is byte-safe since ReduceOp is TrivialRegisterPassable):

  • Small (<= _WARP_SHUFFLE_MAX_WORDS uint32 words): a register-only within-warp warp-shuffle butterfly, then a sparse one-entry-per-warp shmem store + a second shuffle over warp 0 (mirrors legacy welford_block_all_reduce) β€” BLOCK_SIZE/WARP_SIZE stores + one barrier, ~30x less shmem traffic than the tree.
  • Large: a log2(BLOCK_SIZE)-step block-wide shmem tree. Past a few words the per-word shuffle + register pressure makes the tree comparable or better.

Parameters:

  • ​S (ReduceOp): The monoid type being combined.

Args:

  • ​state (S): The per-thread state; on return, holds the block-wide combined value on every thread.