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

map_reduce

def map_reduce[simd_width: SIMDLength, dtype: DType, acc_type: DType, origins_gen: OriginSet, input_gen_fn: def[dtype: DType, width: Int](Int) capturing thin -> SIMD[dtype, width], origins_vec: OriginSet, reduce_vec_to_vec_fn: def[acc_type: DType, dtype: DType, width: SIMDLength](SIMD[acc_type, width], SIMD[dtype, width]) capturing thin -> SIMD[acc_type, width], reduce_vec_to_scalar_fn: def[dtype: DType, width: SIMDLength](SIMD[dtype, width]) thin -> Scalar[dtype]](dst: Span[Scalar[dtype]], init: Scalar[acc_type]) -> Scalar[acc_type]

Stores the result of calling input_gen_fn in dst and simultaneously reduce the result using a custom reduction function.

Parameters:

simd_width (SIMDLength): The vector width for the computation.
dtype (DType): The buffer elements dtype.
acc_type (DType): The dtype of the reduction accumulator.
origins_gen (OriginSet): The OriginSet of captured arguments by the input_gen_fn.
input_gen_fn (def[dtype: DType, width: Int](Int) capturing thin -> SIMD[dtype, width]): A function that generates inputs to reduce.
origins_vec (OriginSet): The OriginSet of captured arguments by the reduce_vec_to_vec_fn.
reduce_vec_to_vec_fn (def[acc_type: DType, dtype: DType, width: SIMDLength](SIMD[acc_type, width], SIMD[dtype, width]) capturing thin -> SIMD[acc_type, width]): A mapping function. This function is used to combine (accumulate) two chunks of input data: e.g. we load two 8xfloat32 vectors of elements and need to reduce them into a single 8xfloat32 vector.
reduce_vec_to_scalar_fn (def[dtype: DType, width: SIMDLength](SIMD[dtype, width]) thin -> Scalar[dtype]): A reduction function. This function is used to reduce a vector to a scalar. E.g. when we got 8xfloat32 vector and want to reduce it to an float32 scalar.

Args:

dst (Span[Scalar[dtype]]): The output buffer.
init (Scalar[acc_type]): The initial value to use in accumulator.

Returns:

Scalar[acc_type]: The computed reduction value.

def map_reduce[simd_width: SIMDLength, dtype: DType, acc_type: DType, origins_gen: OriginSet, input_gen_fn: def[dtype: DType, width: Int](Int) capturing thin -> SIMD[dtype, width], origins_vec: OriginSet, reduce_vec_to_vec_fn: def[acc_type: DType, dtype: DType, width: SIMDLength](SIMD[acc_type, width], SIMD[dtype, width]) capturing thin -> SIMD[acc_type, width], reduce_vec_to_scalar_fn: def[dtype: DType, width: SIMDLength](SIMD[dtype, width]) thin -> Scalar[dtype], output_fn: def[dtype_: DType, width: SIMDLength, alignment: Int](idx: Int, val: SIMD[dtype_, width]) capturing thin -> None](length: Int, init: Scalar[acc_type]) -> Scalar[acc_type]

Performs a vectorized map-reduce operation over a sequence.

Parameters:

simd_width (SIMDLength): The SIMD vector width to use.
dtype (DType): The data type of the input elements.
acc_type (DType): The data type of the accumulator.
origins_gen (OriginSet): Origin set for the input generation function.
input_gen_fn (def[dtype: DType, width: Int](Int) capturing thin -> SIMD[dtype, width]): Function that generates input values at each index.
origins_vec (OriginSet): Origin set for the reduction function.
reduce_vec_to_vec_fn (def[acc_type: DType, dtype: DType, width: SIMDLength](SIMD[acc_type, width], SIMD[dtype, width]) capturing thin -> SIMD[acc_type, width]): Function that reduces a vector into the accumulator.
reduce_vec_to_scalar_fn (def[dtype: DType, width: SIMDLength](SIMD[dtype, width]) thin -> Scalar[dtype]): Function that reduces a final vector to a scalar.
output_fn (def[dtype_: DType, width: SIMDLength, alignment: Int](idx: Int, val: SIMD[dtype_, width]) capturing thin -> None): Function to output intermediate results.

Args:

length (Int): The number of elements to process.
init (Scalar[acc_type]): The initial accumulator value.

Returns:

Scalar[acc_type]: The final reduced scalar value.