Accumulate / Prefix Sum / Scan
AcceleratedKernels.accumulate! — Function
accumulate!(op, v::AbstractArray; kwargs...) -> v
accumulate!(op, dst::AbstractArray, src::AbstractArray; kwargs...) -> dst
accumulate!(
op, dst, src;
backend=nothing,
init=<none>,
neutral=nothing,
dims::Union{Nothing, Integer}=nothing,
inclusive::Bool=true,
acctype=nothing,
alg::Algorithm=Auto(),
workspace=nothing,
)Compute accumulated running totals along a sequence by applying a binary operator to all elements up to the current one; often used in GPU programming as a first step in finding / extracting subsets of data. The first form scans v in place; the second writes the scan of src to dst, which must be src itself or not overlap it.
Other names: prefix sum, thrust::scan, cumulative sum; inclusive (or exclusive) if the first element is included in the accumulation (or not).
With dims=nothing the array is scanned in linear order, whatever its number of dimensions; with an integer dims, each slice along that dimension is scanned independently. A dims beyond the array's dimensions makes every slice one element long.
The contract, which is AcceleratedKernels' own (see Differences from Base):
- Algebra.
opmust be associative, as elements are combined in parallel; it does not need to be commutative, as every combination keeps the elements in order (matrix products are fine). Non-associative operators such as-are not supported. - Inclusive scans (the default): with running values
y,y[1] = Base.reduce_first(op, src[1])withoutinit, andop(init, src[1])with it; theny[k] = op(y[k-1], src[k]), anddst[k] = y[k].initis applied once per slice. - Exclusive scans (
inclusive=false):y[1]isinit, or the neutral element ofopwithout it (anArgumentErrorwhen none is known); theny[k] = op(y[k-1], src[k-1]). - Neutral element.
neutral(a two-sided identity ofop) seeds partial results; it defaults toGPUArraysCore.neutral_element(op, T)where that is defined. For other operators, partial results start from their first element instead, so inclusive scans need no neutral element. - Running-value type. The running values have one type, and are converted to
eltype(dst)only when stored:acctypewhen it is given, else the type the fold ofopsettles on, starting fromeltype(dst)(joined withinit's type) and the elements, as formapreducedim!(e.g.Intforadd_sumoverInt8). So with the usual operators and noacctype, a scan does not run in a type narrower than its destination:Float32s scanned into aFloat64array are summed inFloat64. Where the running type differs fromeltype(dst), the scan runs in a scratch array of that type. Anacctypethat cannot hold the running values at all is anArgumentError; whether their values fit is the caller's obligation.accumulateallocates the fold type frominit's type and the elements.
alg is Auto() by default: CPUThreads.Partitioned on the host, and on GPUs ScanPrefixes for whole arrays and SliceScan along dims, with the device's settings. DecoupledLookback is available on backends that support it. backend is derived from dst and src.
On the host, accumulation is typically a memory-bound operation, so multithreaded accumulation only becomes faster for more compute-heavy operations that hide memory latency, e.g. accumulating tuples or structs, or expensive operators.
workspace takes the scratch memory of a workspace made for the same call, so that the scan allocates none of its own.
Examples
Example computing an inclusive prefix sum (the typical GPU "scan"):
import AcceleratedKernels as AK
using oneAPI
v = oneAPI.ones(Int32, 100_000)
AK.accumulate!(+, v) # 1, 2, 3, ...
AK.accumulate!(+, v; inclusive=false) # 0, 1, 3, ...: starts from the neutral element
# Choose the algorithm and its settings
AK.accumulate!(+, v; alg=AK.ScanPrefixes(block_size=512))AcceleratedKernels.accumulate — Function
accumulate(op, v::AbstractArray; init=<none>, kwargs...)Out-of-place version of accumulate!, with the same keywords. The result's element type is acctype when it is given, else the type the fold of op settles on from init's type (when given) and the elements, e.g. an Int array for accumulate(+, Int8[1, 2]; init=0) and for accumulate(Base.add_sum, Int8[1, 2]).
Auto(), the default alg, uses CPUThreads.Partitioned on the host; on GPUs, ScanPrefixes for whole arrays and SliceScan along dims, with the device's settings. Pass an algorithm to choose it and its settings yourself:
AcceleratedKernels.ScanAlgorithm — Type
ScanAlgorithm <: AlgorithmSupertype of the scan (accumulate) algorithms: ScanPrefixes, DecoupledLookback and SliceScan. Scans also accept CPUThreads.Partitioned.
AcceleratedKernels.ScanPrefixes — Type
ScanPrefixes(; block_size=nothing, items_per_thread=nothing)GPU scan of a whole array: each block scans a tile of block_size * items_per_thread elements (block_size a power of two up to 1024), then the tiles' totals are scanned and added to later tiles. No block waits for another, so it runs on every backend. The default items_per_thread is at most 8, fewer for wide element types, to keep the tile in local memory.
AcceleratedKernels.DecoupledLookback — Type
DecoupledLookback(; block_size=nothing, items_per_thread=nothing)GPU scan of a whole array in which each block looks back at earlier blocks' published prefixes instead of waiting for a separate pass over the tiles' totals; the tiles are those of ScanPrefixes. It needs device-scope memory ordering and forward progress between blocks, so it is rejected on backends that do not guarantee them (currently all but CUDA and AMDGPU). Elements keep their order, but how they are grouped depends on which blocks have finished, so floating-point results can differ between runs.
AcceleratedKernels.SliceScan — Type
SliceScan(; block_size=nothing)GPU scan of each slice along dims: one thread per slice when there are more slices than elements per slice, otherwise one block (block_size a power of two up to 1024) per slice. The algorithm for scans along dims.
v = CuArray(rand(Int32(1):Int32(100), 1_000_000))
AK.accumulate!(+, v; init=Int32(0), alg=AK.ScanPrefixes(block_size=512)) # 512 threads, 8 items per thread
AK.accumulate!(+, v; init=Int32(0), alg=AK.DecoupledLookback()) # CUDA and AMDGPU only
m = CuArray(rand(Int32(1):Int32(100), 100, 10_000))
AK.accumulate(+, m; init=Int32(0), dims=2, alg=AK.SliceScan(block_size=128))