GitHub

Find All / Stream Compaction

AcceleratedKernels.findall — Function
findall(A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), workspace=nothing)
findall(pred, A::AbstractArray; items=keys(A), backend=nothing, alg=Auto(), workspace=nothing)

Stream compaction: select, in order, the elements of items at the positions where A is true, or where pred returns true for A's elements. Values used as conditions, and pred's results, must be Bool. The result is a new vector of eltype(items) on the backend.

Positions are ordinal: the k-th element of A, in the order of eachindex(A), selects the k-th element of items, which may be any array with A's length (so offset axes are no problem). What is selected is the caller's choice:

  • keys(A) (the default) gives A's indices: Ints for a vector, CartesianIndexes for other arrays, including 0-dimensional ones.
  • LinearIndices(A) gives linear indices of any array.
  • An array of values selects those values: AK.findall(mask; items=A) is A[mask] for a mask of A's shape, in one pass instead of findall and a gather.

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

alg is Auto() by default: CPUThreads.Partitioned on the host and ScanScatter on GPUs, with the device's settings. backend is derived from A and items. workspace takes the scratch memory of a workspace made for the same call (the mask of the predicate form and the counts), so that findall allocates only its result.

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]
AK.findall(x -> x > 0, v; items=v)      # Int32[5, 8, 3]

m = CUDA.CuArray(Bool[1 0; 0 1])
AK.findall(m)                           # [CartesianIndex(1, 1), CartesianIndex(2, 2)]
AK.findall(m; items=LinearIndices(m))   # [1, 4]
source

findall selects items: indices by default, or anything else of the input's length, such as the values themselves.

v = CuArray(rand(Float32, 1000))
AK.findall(x -> x > 0.5f0, v)                  # indices
AK.findall(x -> x > 0.5f0, v; items=v)         # the selected values, as v[v .> 0.5f0]

Auto() uses CPUThreads.Partitioned on the host and ScanScatter on GPUs, with the device's settings.

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

Stable GPU stream compaction: each block counts the selected elements of its tile of block_size * items_per_thread elements (block_size a power of two up to 1024), the counts are scanned, and a second pass scatters the selected indices.

source