Mojo function

map_reduce

map_reduce[simd_width: Int, size: Dim, dtype: DType, acc_type: DType, origins_gen: OriginSet, input_gen_fn: fn[dtype: DType, width: Int](Int) capturing -> SIMD[dtype, width], origins_vec: OriginSet, reduce_vec_to_vec_fn: fn[acc_type: DType, dtype: DType, width: Int](SIMD[acc_type, width], SIMD[dtype, width]) capturing -> SIMD[acc_type, width], reduce_vec_to_scalar_fn: fn[dtype: DType, width: Int](SIMD[dtype, width]) -> Scalar[dtype]](dst: NDBuffer[dtype, 1, origin, size], 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 (Int): The vector width for the computation.
size (Dim): The buffer size.
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 (fn[dtype: DType, width: Int](Int) capturing -> 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 (fn[acc_type: DType, dtype: DType, width: Int](SIMD[acc_type, width], SIMD[dtype, width]) capturing -> 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 (fn[dtype: DType, width: Int](SIMD[dtype, width]) -> 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 (NDBuffer): The output buffer.
init (Scalar): The initial value to use in accumulator.

Returns:

Scalar: The computed reduction value.

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

Performs a vectorized map-reduce operation over a sequence.

Parameters:

simd_width (Int): 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 (fn[dtype: DType, width: Int](Int) capturing -> 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 (fn[acc_type: DType, dtype: DType, width: Int](SIMD[acc_type, width], SIMD[dtype, width]) capturing -> SIMD[acc_type, width]): Function that reduces a vector into the accumulator.
reduce_vec_to_scalar_fn (fn[dtype: DType, width: Int](SIMD[dtype, width]) -> Scalar[dtype]): Function that reduces a final vector to a scalar.
output_fn (fn[dtype_: DType, width: Int, alignment: Int](idx: Int, val: SIMD[dtype_, width]) capturing -> None): Function to output intermediate results.

Args:

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

Returns:

Scalar: The final reduced scalar value.