diff --git a/README.md b/README.md index b01c30aa..b0234cdf 100644 --- a/README.md +++ b/README.md @@ -281,6 +281,7 @@ If you need other algorithms in your work that may be of general use, please ope | [Accumulation](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/accumulate/) | `accumulate` `accumulate!` | `prefix_sum` `thrust::scan` `cumsum` | | [Binary Search](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/binarysearch/) | `searchsortedfirst` `searchsortedfirst!` | `std::lower_bound` | | | `searchsortedlast` `searchsortedlast!` | `thrust::upper_bound` | +| [Find All](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/findall/) | `findall` | `thrust::copy_if` `cub::DeviceSelect` `nonzero` | | [Predicates](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/predicates/) | `all` `any` | | | [Arithmetics](https://juliagpu.github.io/AcceleratedKernels.jl/stable/api/arithmetics/) | `sum` `prod` `minimum` `maximum` `count` `cumsum` `cumprod` | | diff --git a/docs/make.jl b/docs/make.jl index 58b4632e..3361e3d3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -28,6 +28,7 @@ makedocs(; "MapReduce" => "api/mapreduce.md", "Accumulate" => "api/accumulate.md", "Binary Search" => "api/binarysearch.md", + "Find All" => "api/findall.md", "Predicates" => "api/predicates.md", "Arithmetics" => "api/arithmetics.md", "Custom Structs" => "api/custom_structs.md", diff --git a/docs/src/api/findall.md b/docs/src/api/findall.md new file mode 100644 index 00000000..5f2aa606 --- /dev/null +++ b/docs/src/api/findall.md @@ -0,0 +1,6 @@ +### Find All / Stream Compaction + +```@docs +AcceleratedKernels.findall +AcceleratedKernels.ScanScatter +``` diff --git a/src/AcceleratedKernels.jl b/src/AcceleratedKernels.jl index 35d12603..c0205561 100644 --- a/src/AcceleratedKernels.jl +++ b/src/AcceleratedKernels.jl @@ -32,6 +32,7 @@ include("map.jl") include("sort/sort.jl") include("reduce/reduce.jl") include("accumulate/accumulate.jl") +include("findall.jl") include("reverse.jl") include("searchsorted.jl") include("predicates.jl") diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index f4d1a481..98d33bbf 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -15,6 +15,49 @@ end function _decoupled_fence end +# Exclusive scan of one value per thread in local memory. All threads in the block must call it. +@inline function block_exclusive_scan!(@context, op, totals, seed, block_size, ithread) + # Up-sweep. Use index-sized counters for block sizes of 256 or more. + offset = one(ithread) + d = block_size >> 0x1 + while d > 0x0 + @synchronize() + if ithread < d + ai = offset * (0x2 * ithread + 0x1) - 0x1 + bi = offset * (0x2 * ithread + 0x2) - 0x1 + totals[bi + 0x1] = op(totals[bi + 0x1], totals[ai + 0x1]) + end + offset = offset << 0x1 + d = d >> 0x1 + end + + @synchronize() + block_total = op(seed, totals[block_size]) + @synchronize() + if ithread == 0x0 + totals[block_size] = seed + end + + # Down-sweep to an exclusive scan. + d = one(ithread) + while d < block_size + offset = offset >> 0x1 + @synchronize() + if ithread < d + ai = offset * (0x2 * ithread + 0x1) - 0x1 + bi = offset * (0x2 * ithread + 0x2) - 0x1 + t = totals[ai + 0x1] + totals[ai + 0x1] = totals[bi + 0x1] + totals[bi + 0x1] = op(totals[bi + 0x1], t) + end + d = d << 0x1 + end + @synchronize() + + return totals[ithread + 0x1], block_total +end + + # Register-raking block scan with striped loads and stores. @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_block!( op, v, init, neutral, @@ -53,50 +96,13 @@ function _decoupled_fence end k += 1 end thread_totals[ithread + 0x1] = acc - @synchronize() # Scan the per-thread totals. Later blocks receive their carry from the # second kernel. seed = iblock == 0x0 ? init : neutral - - # Use index-sized counters for block sizes of 256 or more. - offset = one(ithread) - d = block_size >> 0x1 - while d > 0x0 - @synchronize() - if ithread < d - ai = offset * (0x2 * ithread + 0x1) - 0x1 - bi = offset * (0x2 * ithread + 0x2) - 0x1 - thread_totals[bi + 0x1] = - op(thread_totals[bi + 0x1], thread_totals[ai + 0x1]) - end - offset = offset << 0x1 - d = d >> 0x1 - end - - @synchronize() - block_total = op(seed, thread_totals[block_size]) - @synchronize() - if ithread == 0x0 - thread_totals[block_size] = seed - end - - # Down-sweep to an exclusive scan. - d = one(ithread) - while d < block_size - offset = offset >> 0x1 - @synchronize() - if ithread < d - ai = offset * (0x2 * ithread + 0x1) - 0x1 - bi = offset * (0x2 * ithread + 0x2) - 0x1 - t = thread_totals[ai + 0x1] - thread_totals[ai + 0x1] = thread_totals[bi + 0x1] - thread_totals[bi + 0x1] = op(thread_totals[bi + 0x1], t) - end - d = d << 0x1 - end - @synchronize() - thread_prefix = thread_totals[ithread + 0x1] + thread_prefix, block_total = block_exclusive_scan!( + @context, op, thread_totals, seed, block_size, ithread, + ) # DecoupledLookback keeps later blocks inclusive until the carry pass. block_inclusive = inclusive || (iblock != 0x0 && !isnothing(flags)) diff --git a/src/findall.jl b/src/findall.jl new file mode 100644 index 00000000..055a161d --- /dev/null +++ b/src/findall.jl @@ -0,0 +1,305 @@ +abstract type FindallAlgorithm end + +""" + 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. +""" +Base.@kwdef struct ScanScatter <: FindallAlgorithm + block_size::Int = 256 + items_per_thread::Int = 16 +end + + +findall_algorithm(alg::ScanScatter) = alg +function findall_algorithm(alg::FindallAlgorithm) + throw(ArgumentError("$(typeof(alg)) is not supported by findall")) +end + + +@inline findall_index(indices::AbstractUnitRange, position) = + first(indices) + position - 1 +@inline findall_index(indices::LinearIndices{1}, position) = + first(indices) + position - 1 +@inline findall_index(indices, position) = @inbounds indices[position] + + +# With `out === nothing`, compute block counts. Otherwise, `block_counts` contains their +# inclusive prefix scan and the kernel scatters the selected indices. +@kernel cpu=false inbounds=true unsafe_indices=true function findall_block!( + out, @Const(bools), block_counts, input_indices, output_indices, ::Val{ITEMS}, +) where ITEMS + @uniform block_size = @groupsize()[1] + tile = @localmem UInt8 (block_size * ITEMS,) + thread_counts = @localmem Int (block_size,) + + len = length(bools) + iblock = @index(Group, Linear) - 0x1 + ithread = @index(Local, Linear) - 0x1 + block_offset = iblock * block_size * ITEMS + + j = 0 + while j < ITEMS + p = j * block_size + ithread + position = block_offset + p + tile[p + 0x1] = position < len ? + UInt8(bools[findall_index(input_indices, position + 0x1)]) : 0x0 + j += 1 + end + @synchronize() + + run = ithread * ITEMS + count = 0 + k = 0 + while k < ITEMS + count += tile[run + k + 0x1] + k += 1 + end + thread_counts[ithread + 0x1] = count + + seed = (isnothing(out) || iblock == 0x0) ? 0 : block_counts[iblock] + pos, block_total = block_exclusive_scan!( + @context, +, thread_counts, seed, block_size, ithread, + ) + + if isnothing(out) + if ithread == 0x0 + block_counts[iblock + 0x1] = block_total + end + else + k = 0 + while k < ITEMS + if tile[run + k + 0x1] != 0x0 + pos += 1 + position = block_offset + run + k + 0x1 + out[pos] = findall_index(output_indices, position) + end + k += 1 + end + end +end + + +function findall_temp(bools, backend, len, temp) + if isnothing(temp) + return KernelAbstractions.allocate(backend, Int, len) + end + + @argcheck get_backend(temp) === backend + @argcheck eltype(temp) === Int + @argcheck length(temp) >= len + @argcheck !Base.mightalias(temp, bools) + view(temp, 1:len) +end + + +function findall_gpu( + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg::ScanScatter; + temp, +) where I + block_size = alg.block_size + items_per_thread = alg.items_per_thread + @argcheck 1 <= block_size <= 1024 + @argcheck ispow2(block_size) + @argcheck items_per_thread > 0 + + isempty(bools) && return similar(bools, I, 0) + + elems_per_block = block_size * items_per_thread + num_blocks = cld(length(bools), elems_per_block) + block_counts = findall_temp(bools, backend, num_blocks, temp) + input_indices = eachindex(bools) + items = Val(items_per_thread) + + kernel! = findall_block!(backend, block_size) + kernel!(nothing, bools, block_counts, input_indices, output_indices, items; + ndrange=num_blocks * block_size) + accumulate!(+, block_counts, backend; init=0) + n = @allowscalar block_counts[end] + + out = similar(bools, I, n) + if n > 0 + kernel!(out, bools, block_counts, input_indices, output_indices, items; + ndrange=num_blocks * block_size) + end + out +end + + +function findall_cpu( + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend; + max_tasks::Int, + min_elems::Int, + temp, +) where I + input_indices = eachindex(bools) + tp = TaskPartitioner(length(bools), max_tasks, min_elems) + if tp.num_tasks == 1 + out = similar(bools, I, Base.count(bools)) + findall_section!(out, bools, input_indices, output_indices, Base.OneTo(length(bools)), 0) + return out + end + + task_counts = findall_temp(bools, backend, tp.num_tasks, temp) + itask_partition(tp) do itask, positions + task_counts[itask] = Base.count( + position -> @inbounds(bools[findall_index(input_indices, position)]), positions, + ) + end + cumsum!(task_counts, task_counts) + + out = similar(bools, I, task_counts[end]) + itask_partition(tp) do itask, positions + offset = itask == 1 ? 0 : task_counts[itask - 1] + findall_section!(out, bools, input_indices, output_indices, positions, offset) + end + out +end + + +function findall_section!(out, bools, input_indices, output_indices, positions, pos) + @inbounds for position in positions + if bools[findall_index(input_indices, position)] + pos += 1 + out[pos] = findall_index(output_indices, position) + end + end + out +end + + +function findall_impl( + bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend; + alg::FindallAlgorithm=ScanScatter(), + max_tasks::Int=Threads.nthreads(), + min_elems::Int=1, + prefer_threads::Bool=true, + temp::Union{Nothing, AbstractArray}=nothing, +) where I + alg = findall_algorithm(alg) + + if use_gpu_algorithm(backend, prefer_threads) + findall_gpu(bools, I, output_indices, backend, alg; temp) + else + findall_cpu(bools, I, output_indices, backend; max_tasks, min_elems, temp) + end +end + + +findall_output_indices(v, ::Type{Int}) = LinearIndices(v) +findall_output_indices(v, ::Type) = CartesianIndices(axes(v)) + + +function findall_bools( + pred, v::AbstractArray, backend::Backend, temp_bools; + max_tasks, min_elems, prefer_threads, block_size, +) + bools = if isnothing(temp_bools) + similar(v, Bool) + else + @argcheck get_backend(temp_bools) === backend + @argcheck eltype(temp_bools) === Bool + @argcheck axes(temp_bools) == axes(v) + @argcheck !Base.mightalias(temp_bools, v) + temp_bools + end + input_indices = eachindex(v) + bool_indices = eachindex(bools) + foreachindex(Base.OneTo(length(v)), backend; + max_tasks, min_elems, prefer_threads, block_size) do position + input_index = findall_index(input_indices, position) + bool_index = findall_index(bool_indices, position) + @inbounds bools[bool_index] = pred(v[input_index]) ? true : false + end + bools +end + + +""" + 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 +```julia +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)] +``` +""" +function findall( + values::AbstractArray, backend::Backend=get_backend(values); + 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, +) + alg = findall_algorithm(alg) + bools = if values isa AbstractArray{Bool} + isnothing(temp_bools) || + throw(ArgumentError("temp_bools is not used for a Bool mask")) + values + else + findall_bools(identity, values, backend, temp_bools; + max_tasks, min_elems, prefer_threads, block_size=alg.block_size) + end + I = keytype(values) + output_indices = findall_output_indices(values, I) + findall_impl(bools, I, output_indices, backend; + alg, max_tasks, min_elems, prefer_threads, temp) +end + + +function findall( + pred, v::AbstractArray, backend::Backend=get_backend(v); + 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, +) + alg = findall_algorithm(alg) + bools = findall_bools(pred, v, backend, temp_bools; + max_tasks, min_elems, prefer_threads, block_size=alg.block_size) + I = ndims(v) == 0 ? Int : keytype(v) + output_indices = findall_output_indices(v, I) + findall_impl(bools, I, output_indices, backend; + alg, max_tasks, min_elems, prefer_threads, temp) +end diff --git a/test/generic/findall.jl b/test/generic/findall.jl new file mode 100644 index 00000000..3adbdb25 --- /dev/null +++ b/test/generic/findall.jl @@ -0,0 +1,174 @@ +struct FindallCallable end +(::FindallCallable)(x) = x > zero(x) + +struct OtherFindallAlgorithm <: AK.FindallAlgorithm end + +struct FindallOffsetVector{T, V <: AbstractVector{T}} <: AbstractVector{T} + data::V + offset::Int +end + +Base.size(v::FindallOffsetVector) = size(v.data) +Base.axes(v::FindallOffsetVector) = + (Base.IdentityUnitRange((firstindex(v.data) + v.offset):(lastindex(v.data) + v.offset)),) +Base.IndexStyle(::Type{<:FindallOffsetVector}) = IndexLinear() +Base.getindex(v::FindallOffsetVector, i::Int) = v.data[i - v.offset] +Base.setindex!(v::FindallOffsetVector, x, i::Int) = (v.data[i - v.offset] = x) +Base.similar(v::FindallOffsetVector, ::Type{T}) where T = + FindallOffsetVector(similar(v.data, T), v.offset) +Base.similar(v::FindallOffsetVector, ::Type{T}, dims::Dims) where T = similar(v.data, T, dims) + + +@testset "findall" begin + Random.seed!(0) + + default_alg = AK.ScanScatter() + tile_size = default_alg.block_size * default_alg.items_per_thread + edge_sizes = [0, 1, 2, 3, tile_size - 1, tile_size, tile_size + 1, + 2tile_size - 1, 2tile_size, 2tile_size + 1, 10_000] + test_types = valid_backend_eltypes(BACKEND, (Int32, Float32, Float64)) + + @testset "predicate" begin + for T in test_types, n in edge_sizes + pred = T <: AbstractFloat ? (x -> x > oftype(x, 0.5)) : (x -> x > zero(x)) + h = T <: AbstractFloat ? rand(T, n) : rand(T(-5):T(5), n) + v = array_from_host(h) + @test Array(AK.findall(pred, v; prefer_threads)) == findall(pred, h) + end + + h = collect(Int32, -10:10) + v = array_from_host(h) + @test Array(AK.findall(FindallCallable(), v; prefer_threads)) == findall(x -> x > 0, h) + + if prefer_threads + calls = Ref(0) + pred = x -> (calls[] += 1; isodd(x)) + h = collect(1:100) + @test AK.findall(pred, h; max_tasks=1) == findall(isodd, h) + @test calls[] == length(h) + @test_throws TypeError AK.findall(Returns(1), [1]; max_tasks=1) + end + end + + @testset "mask" begin + for n in edge_sizes + h = rand(Bool, n) + v = array_from_host(h) + out = AK.findall(v; prefer_threads) + @test Array(out) == findall(h) + @test eltype(out) == Int + end + + if prefer_threads + @test AK.findall(Any[true, false, true]) == findall(Any[true, false, true]) + @test AK.findall(Any[true, false, true]; temp_bools=Vector{Bool}(undef, 3)) == + findall(Any[true, false, true]) + @test_throws TypeError AK.findall([1]) + @test_throws TypeError AK.findall(Any[true, missing]; max_tasks=1) + + scalar = Array{Any}(undef) + scalar[] = true + @test AK.findall(scalar) == findall(scalar) + end + end + + @testset "dimensions and keys" begin + for shape in ([4, 2], [1, 6], [64, 64], [8, 8, 8]) + h = rand(Float32, shape...) + v = array_from_host(h) + out = AK.findall(x -> x > 0.5f0, v; prefer_threads) + @test Array(out) == findall(x -> x > 0.5f0, h) + @test eltype(out) == CartesianIndex{length(shape)} + + hb = rand(Bool, shape...) + @test Array(AK.findall(array_from_host(hb); prefer_threads)) == findall(hb) + end + + for value in (false, true) + h = fill(value) + @test Array(AK.findall(array_from_host(h); prefer_threads)) == findall(h) + end + for value in (0.25f0, 0.75f0) + h = fill(value) + pred = x -> x > 0.5f0 + @test Array(AK.findall(pred, array_from_host(h); prefer_threads)) == findall(pred, h) + end + + if prefer_threads + h = FindallOffsetVector([-1, 1, -2, 2, 0], -3) + mask = FindallOffsetVector(Bool[false, true, true, false, true], -3) + for max_tasks in (1, 4) + @test AK.findall(x -> x > 0, h, BACKEND; max_tasks, min_elems=1) == + [-1, 1] + @test AK.findall(mask, BACKEND; max_tasks, min_elems=1) == findall(mask) + end + + h = collect(1:20) + v = @view h[2:2:20] + @test AK.findall(isodd, v) == findall(isodd, v) + @test AK.findall(iszero, reshape(Int[], 0, 2)) == + findall(iszero, reshape(Int[], 0, 2)) + end + end + + @testset "selection extremes" begin + for n in (0, 1, 2, 1000), h in (trues(n), falses(n)) + values = collect(h) + @test Array(AK.findall(array_from_host(values); prefer_threads)) == findall(values) + end + + v = array_from_host(collect(Int32, 1:1000)) + @test Array(AK.findall(x -> x > 0, v; prefer_threads)) == collect(1:1000) + @test Array(AK.findall(x -> x < 0, v; prefer_threads)) == Int[] + end + + @testset "random sizes" begin + for _ in 1:100 + n = rand(1:100_000) + h = rand(Float32, n) + v = array_from_host(h) + @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads)) == + findall(x -> x > 0.5f0, h) + end + end + + @testset "configuration and buffers" begin + h = rand(Float32, 10_000) + v = array_from_host(h) + for block_size in (32, 64, 128, 256), items_per_thread in (1, 3, 8) + alg = AK.ScanScatter(; block_size, items_per_thread) + @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads, alg)) == + findall(x -> x > 0.5f0, h) + end + + for (max_tasks, min_elems) in ((1, 1), (2, 100), (4, 1000)) + @test Array(AK.findall(x -> x > 0.5f0, v; prefer_threads, max_tasks, min_elems)) == + findall(x -> x > 0.5f0, h) + end + + alg = AK.ScanScatter(block_size=64, items_per_thread=3) + temp = similar(v, Int, max(4, cld(length(v), alg.block_size * alg.items_per_thread))) + temp_bools = similar(v, Bool) + @test Array(AK.findall(x -> x > 0.5f0, v; + prefer_threads, max_tasks=4, alg, temp, temp_bools)) == + findall(x -> x > 0.5f0, h) + + @test_throws ArgumentError AK.findall(v; prefer_threads, alg=OtherFindallAlgorithm()) + @test_throws ArgumentError AK.findall(identity, temp_bools; + prefer_threads, temp_bools) + @test_throws ArgumentError AK.findall(identity, v; prefer_threads, + temp_bools=reshape(similar(v, Bool), :, 1)) + + if !prefer_threads + bools = array_from_host(rand(Bool, length(v))) + @test_throws ArgumentError AK.findall(bools; prefer_threads, + alg=AK.ScanScatter(block_size=192)) + @test_throws ArgumentError AK.findall(bools; prefer_threads, + alg=AK.ScanScatter(items_per_thread=0)) + @test_throws ArgumentError AK.findall(bools; prefer_threads, + temp=similar(v, Int32, 100)) + @test_throws ArgumentError AK.findall(bools; prefer_threads, + temp=similar(v, Int, 1)) + end + end +end