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 function

reduce

def reduce[val_type: DType, simd_width: SIMDLength, FuncType: def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable, //, shuffle: def[dtype: DType, simd_width: SIMDLength](val: SIMD[dtype, simd_width], offset: UInt32) thin -> SIMD[dtype, simd_width]](val: SIMD[val_type, simd_width], func: FuncType) -> SIMD[val_type, simd_width]

Performs a generic warp-wide reduction operation using shuffle operations.

This is a convenience wrapper around lane_group_reduce that operates on the entire warp. It allows customizing both the shuffle operation and reduction function.

Example:

    from max.gpu.primitives.warp import reduce, shuffle_down

    # Compute warp-wide sum using shuffle down
    def add[dtype: DType, width: SIMDLength](x: SIMD[dtype, width], y: SIMD[dtype, width]) -> SIMD[dtype, width]:
        return x + y

    val = SIMD[.float32, 4](2.0, 4.0, 6.0, 8.0)
    result = reduce[shuffle_down](val, add)

Parameters:

• ​val_type (DType): The data type of the SIMD elements (e.g. float32, int32).
• ​simd_width (SIMDLength): The number of elements in the SIMD vector.
• ​FuncType (def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) -> SIMD[dtype, width] & RegisterPassable & ImplicitlyCopyable): The type of the reduction closure.
• ​shuffle (def[dtype: DType, simd_width: SIMDLength](val: SIMD[dtype, simd_width], offset: UInt32) thin -> SIMD[dtype, simd_width]): A function that performs the warp shuffle operation. Takes a SIMD value and offset and returns the shuffled result.

Args:

• ​val (SIMD[val_type, simd_width]): The SIMD value to reduce. Each lane contributes its value.
• ​func (FuncType): A binary function that combines two SIMD values during reduction. This defines the reduction operation (e.g. add, max, min).

Returns:

SIMD[val_type, simd_width]: A SIMD value containing the reduction result broadcast to all lanes in the warp.