MapReduce

Equivalent to reduce(op, map(f, iterable)), without saving the intermediate mapped collection; can be used to e.g. split documents into words (map) and count the frequency thereof (reduce).

  • Other names: transform_reduce, some fold implementations include the mapping function too.

AcceleratedKernels.mapreduce — Function
mapreduce(
    f, op, src, srcs::AbstractArray...;
    backend=nothing,
    init=<none>,
    neutral=nothing,
    acctype=nothing,
    dims=:,
    alg::Algorithm=Auto(),
    workspace=nothing,
)

Apply f to each element of src and reduce the results with the binary operator op. With dims=: (the default; nothing is accepted too) the result is a scalar, returned on the host. With dims an integer or a collection of integers, the result is a new array on src's backend with size 1 along those dimensions; mapreducedim! reduces into an existing array instead, and states the contract all reductions follow. In short:

  • op must be associative and commutative.
  • init, when given, is applied exactly once, as op(init, partial), and is the result of an empty reduction. Without init, an empty reduction is an ArgumentError: of a whole array, and along dims of any output whose slice is empty. (sum, prod and count give zero or one instead.)
  • The result has the accumulator type: acctype when it is given, else the type the fold of op settles on from init's type (when given) and the mapped elements; op(init, partial) is converted to it. So AK.sum(Int8[1, 2]) is an Int, and AK.sum(Int8[1 2]; dims=1, init=Int16(0)) an array of Ints. A single element has that type too: AK.reduce((a, b) -> a + b, [true]) === 1. (An empty reduction returns init as it is.)
  • neutral, a two-sided identity of op, only seeds partial results and never appears in the result. It defaults to GPUArraysCore.neutral_element(op, T) where that is defined; for other operators, each partial result starts from its first element instead.

These rules are AcceleratedKernels' own, and differ from Base's in a few places, e.g. for empty reductions without init; see Differences from Base.

Several source arrays with equal axes are reduced without materializing the mapped array (f takes one argument per array); mismatched axes throw a DimensionMismatch. A Broadcasted object is also accepted as the single source. Before Julia 1.12, both are materialized when the reduction runs. backend is derived from all sources.

alg is Auto() by default: CPUThreads.Partitioned on the host and BlockReduce on GPUs, with the device's settings. On the host, multithreading reductions only improves performance for operations that hide the memory latency and thread launch overhead, e.g. reductions of tuples or structs, or expensive operators.

workspace takes the scratch memory of a workspace made for the same call, so that the reduction allocates none of its own (it still allocates its result along dims, and before Julia 1.12 a materialized source).

Examples

Computing a sum of squares, reducing down to a scalar that is copied to host:

import AcceleratedKernels as AK
using CUDA

v = CuArray{Int16}(rand(1:1000, 100_000))
vsumsq = AK.mapreduce(x -> Int(x) * x, +, v)

Computing dimensionwise sums of squares in a 2D matrix:

import AcceleratedKernels as AK
using Metal

f(x) = x * x
m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000))
mrowsumsq = AK.mapreduce(f, +, m; dims=1)

Computing a two-input dimensional reduction:

rows = AK.mapreduce((x, y) -> x * y, +, a, b; dims=1)
source

To reduce into an existing array, for example to accumulate into it across calls or to avoid an allocation, use mapreducedim!. Its docstring states the contract every reduction follows: the operator algebra, neutral elements, the accumulator type, how init and the destination's values combine with the result, and empty slices.

AcceleratedKernels.mapreducedim! — Function
mapreducedim!(
    f, op, R::AbstractArray, src;
    backend=get_backend(R),
    init=<none>,
    neutral=nothing,
    overwrite::Bool=false,
    acctype=nothing,
    alg::Algorithm=Auto(),
    workspace=nothing,
) -> R

Reduce src (an array or a Broadcasted object) into R: the dimensions along which R has size 1 and src does not are reduced, and the others must match (R may leave off trailing dimensions of size 1). R must not alias src.

This is the contract of every AcceleratedKernels reduction, modelled on CUB's rather than on Base's (see Differences from Base):

  • Algebra. op must be associative and commutative. Like every GPU reduction, AK combines elements in an order that depends on the algorithm, its settings and the shape, not in element order; for a given setting and shape the order is fixed, so results are reproducible on one device. Non-commutative operators such as string concatenation or matrix products are not supported.
  • Partial results start from neutral, which must be a two-sided identity of op and never appears in the result. It defaults to GPUArraysCore.neutral_element(op, T) where that is defined; otherwise each partial result starts from its first element, Base.mapreduce_first(f, op, x), so no neutral element is needed.
  • Accumulator type. Partial results have one type, and are converted to eltype(R) only when stored. It is acctype when that is given; otherwise the type the fold of op settles on, starting from eltype(R) and the mapped elements. Every element is also a one-element partial result, and partial results are combined with each other, so those types are joined in as well (R's values and init are applied once, and are not partial results). A sum of Int8s accumulates in Int, and Float32s reduced into a Float64 array accumulate in Float64. An acctype that cannot hold the partial results at all is an ArgumentError; whether their values fit is the caller's obligation (a value that does not fit throws an InexactError, where the backend reports errors thrown in kernels).
  • Result. For each output whose slice is not empty, with partial its reduction: R[i] = op(init, partial) when init is given (applied once); else partial with overwrite=true; else op(R[i], partial), folding into R's previous value as Base.mapreducedim! does.
  • Empty slices. An output whose slice is empty is set to init when that is given, and is otherwise not written, with or without overwrite.
  • Errors. Mismatched shapes are a DimensionMismatch. An R that aliases src, an op that inference shows always throws for these types (when there is something to combine; an explicit acctype that op cannot combine at all is rejected whatever the input), a non-bits accumulator type for a kernel algorithm, and an algorithm that cannot run on the backend are ArgumentErrors.

workspace takes the scratch memory of a workspace made for the same call.

import AcceleratedKernels as AK
using CUDA

A = CuArray([1 3; 2 4])
R = CuArray([10 20])
AK.mapreducedim!(identity, +, R, A)                  # [13 27]
AK.mapreducedim!(identity, +, R, A; overwrite=true)  # [3 7]
AK.mapreducedim!(identity, +, R, A; init=100)        # [103 107]
source