GitHub

Reductions

Apply a custom binary operator reduction on all elements in an iterable; can be used to compute minima, sums, counts, etc.

  • Other names: Kokkos:parallel_reduce, fold, aggregate.
AcceleratedKernels.reduce — Function
reduce(op, src::AbstractArray; kwargs...)

Reduce src with the binary operator op, which must be associative and commutative; the keywords are those of mapreduce. Equivalent to mapreduce(identity, op, src; kwargs...).

Examples

import AcceleratedKernels as AK
using Metal

m = MtlArray(rand(Int32(1):Int32(100), 10, 100_000))
AK.reduce(+, m)                                     # a host scalar
AK.reduce(max, m; dims=1)                           # a 1×100_000 MtlArray
AK.reduce(+, m; dims=2, alg=AK.BlockReduce(block_size=512))
source

The operator must be associative and commutative: like every GPU reduction, AcceleratedKernels combines elements in an order that depends on the algorithm and its settings, not in element order. init is optional and applied exactly once; without it, an empty reduction is an error, except for sum, prod and count, and operators without a known neutral element need none. Results have the accumulator type, which acctype can set (see mapreducedim! for the contract, and Differences from Base for how it differs from Base's).

AK.sum(CuArray(Int32[]))                # 0, an Int
AK.minimum(CuArray(Int32[]))            # an ArgumentError
AK.reduce((a, b) -> a + b, v)           # no neutral element needed
AK.reduce(+, v; init=10)                # init is added exactly once
AK.sum(CuArray(rand(Float32, 10^6)); acctype=Float64)   # summed in Float64

Auto(), the default alg, uses CPUThreads.Partitioned on the host and BlockReduce on GPUs, with the device's settings. Pass an algorithm to choose its settings yourself:

AcceleratedKernels.BlockReduce — Type
BlockReduce(; block_size=nothing, items_per_thread=nothing, switch_below=nothing)

GPU tree reduction in blocks of block_size threads (a power of two up to 1024), which combine their values in local memory. A whole-array reduction (dims=:) takes at most two launches: in the first, each block reduces tiles of block_size * items_per_thread elements, one after the other, to a partial result, and a single block then reduces those partial results. Reductions to a scalar finish on the host instead once fewer than switch_below values remain; for them, block_size * items_per_thread must be between 2 and typemax(Int32). Reductions along dims choose among several launch shapes by the sizes and strides of the input and output (into a single output, that of whole arrays); items_per_thread and switch_below cannot be set for them, and setting either is an ArgumentError there.

Like every GPU reduction in the Julia ecosystem, it requires an associative and commutative operator: the combination order is fixed for a given setting and shape, but it is not the element order.

source
AcceleratedKernels.CPUThreads.Partitioned — Type
CPUThreads.Partitioned(; max_tasks=nothing, min_elems=nothing)

Split the input into contiguous parts, one per task, process them on Julia threads, and combine the parts' results. Used by reductions, scans, findall and any/all on host arrays. Uses at most max_tasks tasks (default Threads.nthreads()), each with at least min_elems elements (default 1, or the operation's tuning). Only runs on the host backend.

source
m = ROCArray(rand(Float32, 1000, 1000))
AK.reduce(+, m; alg=AK.BlockReduce(block_size=512))            # 512 threads per block
AK.reduce(+, m; alg=AK.BlockReduce(switch_below=1024))         # finish on the host below 1024 values
AK.reduce(+, Array(m); alg=AK.CPUThreads.Partitioned(max_tasks=4))