GitHub

Reductions

Apply a custom binary operator reduction on all elements in an iterable; can be used to compute minima, sums, counts, etc. Also known as Kokkos::parallel_reduce, cub::DeviceReduce, fold or aggregate.

AcceleratedKernels.reduce — Function
reduce(op, src::AbstractArray; kwargs...)

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

Examples

import AcceleratedKernels as AKusing Metalm = MtlArray(rand(Int32(1):Int32(100), 10, 100_000))AK.reduce(+, m)                                     # a host scalarAK.reduce(max, m; dims=1)                           # a 1×100_000 MtlArrayAK.reduce(+, m; dims=2, alg=AK.BlockReduce(block_size=512))
source

The operator must be associative, and by default commutative: as CUB's reductions do, AcceleratedKernels then combines elements in an order that depends on the algorithm and its settings, not in element order. Pass commutative=false for operators such as matrix products, to combine the elements in their 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 IntAK.minimum(CuArray(Int32[]))            # an ArgumentErrorAK.reduce((a, b) -> a + b, v)           # no neutral element neededAK.reduce(+, v; init=10)                # init is added exactly onceAK.reduce(*, matrices; commutative=false)   # products in element orderAK.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 (they finish on the device, and those that read runs of elements per thread take the tuning's items_per_thread).

As CUB's reductions, it requires an associative and commutative operator by default: the combination order is fixed for a given setting and shape, but it is not the element order. With commutative=false (see mapreducedim!), it combines the elements in their 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 blockAK.reduce(+, m; alg=AK.BlockReduce(switch_below=1024))         # finish on the host below 1024 valuesAK.reduce(+, Array(m); alg=AK.CPUThreads.Partitioned(max_tasks=4))