For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

block_reduce

def block_reduce[BLOCK_SIZE: Int, reduce_fn: def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width], dtype: DType, simd_width: SIMDLength](val: SIMD[dtype, simd_width], init: Scalar[dtype]) -> Scalar[dtype]

Performs a block-level reduction of a single SIMD value across all threads in a GPU thread block using warp-level primitives and shared memory.

Parameters:

BLOCK_SIZE (Int): The number of threads per block.
reduce_fn (def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]): The binary reduction function.
dtype (DType): The data type of the elements.
simd_width (SIMDLength): The SIMD vector width.

Args:

val (SIMD[dtype, simd_width]): The per-thread SIMD value to reduce.
init (Scalar[dtype]): The identity value for the reduction.

Returns:

Scalar[dtype]: The reduced scalar result (valid on thread 0).

def block_reduce[BLOCK_SIZE: Int, num_reductions: Int, reduce_fn: def[dtype: DType, width: SIMDLength, reduction_idx: Int](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width], dtype: DType, simd_width: SIMDLength](val: StaticTuple[SIMD[dtype, simd_width], num_reductions], init: StaticTuple[Scalar[dtype], num_reductions]) -> StaticTuple[Scalar[dtype], num_reductions]

Performs a block-level reduction of multiple fused SIMD values across all threads in a GPU thread block using warp shuffles and shared memory.

Parameters:

BLOCK_SIZE (Int): The number of threads per block.
num_reductions (Int): The number of fused reductions to perform.
reduce_fn (def[dtype: DType, width: SIMDLength, reduction_idx: Int](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]): The binary reduction function, parameterized by reduction index.
dtype (DType): The data type of the elements.
simd_width (SIMDLength): The SIMD vector width.

Args:

val (StaticTuple[SIMD[dtype, simd_width], num_reductions]): The per-thread SIMD values to reduce, one per reduction.
init (StaticTuple[Scalar[dtype], num_reductions]): The identity values for each reduction.

Returns:

StaticTuple[Scalar[dtype], num_reductions]: The reduced scalar results (valid on thread 0).