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. op must 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]) without init, and op(init, src[1]) with it; then y[k] = op(y[k-1], src[k]), and dst[k] = y[k]. init is applied once per slice.
  • Exclusive scans (inclusive=false): y[1] is init, or the neutral element of op without it (an ArgumentError when none is known); then y[k] = op(y[k-1], src[k-1]).
  • Neutral element. neutral (a two-sided identity of op) seeds partial results; it defaults to GPUArraysCore.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: acctype when it is given, else the type the fold of op settles on, starting from eltype(dst) (joined with init's type) and the elements, as for mapreducedim! (e.g. Int for add_sum over Int8). So with the usual operators and no acctype, a scan does not run in a type narrower than its destination: Float32s scanned into a Float64 array are summed in Float64. Where the running type differs from eltype(dst), the scan runs in a scratch array of that type. An acctype that cannot hold the running values at all is an ArgumentError; whether their values fit is the caller's obligation. accumulate allocates the fold type from init'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))
source
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]).

source

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.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.

source
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.

source
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.

source
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))