Find All / Stream Compaction

AcceleratedKernels.findallFunction
findall(A::AbstractArray, backend::Backend=get_backend(A);
        alg::FindallAlgorithm=ScanScatter(),
        max_tasks::Int=Threads.nthreads(), min_elems::Int=1,
        prefer_threads::Bool=true,
        temp::Union{Nothing, AbstractArray}=nothing,
        temp_bools::Union{Nothing, AbstractArray}=nothing)
findall(pred, A::AbstractArray, backend::Backend=get_backend(A);
        alg::FindallAlgorithm=ScanScatter(),
        max_tasks::Int=Threads.nthreads(), min_elems::Int=1,
        prefer_threads::Bool=true,
        temp::Union{Nothing, AbstractArray}=nothing,
        temp_bools::Union{Nothing, AbstractArray}=nothing)

Return the indices of the true elements of A, or of the elements for which pred returns true, in the same order and with the same index types as Base.findall. Values used as conditions must be Bool.

The supported inputs are arrays. Dictionaries, other iterables, and scalar inputs accepted by Base.findall are outside the scope of this package.

Settings

  • alg=ScanScatter() selects the GPU algorithm and its tuning parameters.
  • max_tasks=Threads.nthreads() and min_elems=1 control CPU task partitioning.
  • temp=nothing may provide the Int buffer used for block or task counts.
  • temp_bools=nothing may provide the Bool mask for the predicate form or for a mask whose element type is not Bool. It must have the same axes as A and must not alias it.

On a GPU, temp needs at least cld(length(A), alg.block_size * alg.items_per_thread) elements. On a CPU, it needs one element per task used. Omitted buffers are allocated automatically.

Examples

import CUDA
import AcceleratedKernels as AK

v = CUDA.CuArray(Int32[5, -2, 8, -1, 3])
AK.findall(x -> x > 0, v)               # [1, 3, 5]

m = CUDA.CuArray(Bool[1 0; 0 1])
AK.findall(m)                           # [CartesianIndex(1, 1), CartesianIndex(2, 2)]
source
AcceleratedKernels.ScanScatterType
ScanScatter(; block_size=256, items_per_thread=16)

Stable GPU stream compaction using per-block counts, a prefix scan, and a scatter pass. block_size must be a power of two between 1 and 1024; items_per_thread must be positive.

source