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, somefoldimplementations 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:
opmust be associative and commutative.init, when given, is applied exactly once, asop(init, partial), and is the result of an empty reduction. Withoutinit, an empty reduction is anArgumentError: of a whole array, and alongdimsof any output whose slice is empty. (sum,prodandcountgive zero or one instead.)- The result has the accumulator type:
acctypewhen it is given, else the type the fold ofopsettles on frominit's type (when given) and the mapped elements;op(init, partial)is converted to it. SoAK.sum(Int8[1, 2])is anInt, andAK.sum(Int8[1 2]; dims=1, init=Int16(0))an array ofInts. A single element has that type too:AK.reduce((a, b) -> a + b, [true]) === 1. (An empty reduction returnsinitas it is.) neutral, a two-sided identity ofop, only seeds partial results and never appears in the result. It defaults toGPUArraysCore.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)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,
) -> RReduce 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.
opmust 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 ofopand never appears in the result. It defaults toGPUArraysCore.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 isacctypewhen that is given; otherwise the type the fold ofopsettles on, starting fromeltype(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 andinitare applied once, and are not partial results). A sum ofInt8s accumulates inInt, andFloat32s reduced into aFloat64array accumulate inFloat64. Anacctypethat cannot hold the partial results at all is anArgumentError; whether their values fit is the caller's obligation (a value that does not fit throws anInexactError, where the backend reports errors thrown in kernels). - Result. For each output whose slice is not empty, with
partialits reduction:R[i] = op(init, partial)wheninitis given (applied once); elsepartialwithoverwrite=true; elseop(R[i], partial), folding intoR's previous value asBase.mapreducedim!does. - Empty slices. An output whose slice is empty is set to
initwhen that is given, and is otherwise not written, with or withoutoverwrite. - Errors. Mismatched shapes are a
DimensionMismatch. AnRthat aliasessrc, anopthat inference shows always throws for these types (when there is something to combine; an explicitacctypethatopcannot 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 areArgumentErrors.
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]