Skip to main content

/

Docs

Log in

Mojo function

map_reduce

map_reduce[simd_width: Int, size: Dim, type: DType, acc_type: DType, input_gen_fn: fn[DType, Int](Int) capturing -> SIMD[$0, $1], reduce_vec_to_vec_fn: fn[DType, DType, Int](SIMD[$0, $2], SIMD[$1, $2]) capturing -> SIMD[$0, $2], reduce_vec_to_scalar_fn: fn[DType, Int](SIMD[$0, $1]) -> SIMD[$0, 1]](dst: Buffer[type, size, 0], init: SIMD[acc_type, 1]) -> SIMD[$3, 1]

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.
  • ​type (DType): The buffer elements dtype.
  • ​acc_type (DType): The dtype of the reduction accumulator.
  • ​input_gen_fn (fn[DType, Int](Int) capturing -> SIMD[$0, $1]): A function that generates inputs to reduce.
  • ​reduce_vec_to_vec_fn (fn[DType, DType, Int](SIMD[$0, $2], SIMD[$1, $2]) capturing -> SIMD[$0, $2]): 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, Int](SIMD[$0, $1]) -> SIMD[$0, 1]): 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 (Buffer[type, size, 0]): The output buffer.
  • ​init (SIMD[acc_type, 1]): The initial value to use in accumulator.

Returns:

The computed reduction value.

Was this page helpful?