diff --git a/src/findall.jl b/src/findall.jl index e1da08a1..d42a7fd2 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -83,6 +83,28 @@ end @inline findall_index(indices::LinearIndices{1}, position) = first(indices) + position - 1 @inline findall_index(indices, position) = @inbounds indices[firstindex(indices) + position - 1] +# Base converts a linear index to a Cartesian one with checked signed 64-bit divisions, which GPUs +# emulate in long instruction sequences; in the scatter pass, where each thread converts the +# positions of its run in turn, they made small `findall`s of `keys(A)` several times slower than +# those of `LinearIndices(A)`. No axis is empty where a position exists, so unsigned divisions +# without the zero check are exact. +@inline findall_index(::CartesianIndices{0}, position) = CartesianIndex() +@inline function findall_index(indices::CartesianIndices, position) + # (in 32 bits where the indices fit, which some GPUs divide much faster) + offsets = length(indices) <= typemax(UInt32) ? + _cartesian_offsets((position - 1) % UInt32, Base.front(size(indices))) : + _cartesian_offsets((position - 1) % UInt, Base.front(size(indices))) + return CartesianIndex(Base.map((ax, o) -> @inbounds(ax[firstindex(ax) + o % Int]), + indices.indices, offsets)) +end +# The offsets along each dimension of the 0-based linear offset `p` into an array whose leading +# dimensions have the sizes `sizes` +@inline _cartesian_offsets(p::Unsigned, ::Tuple{}) = (p,) +@inline function _cartesian_offsets(p::U, sizes::Tuple) where {U <: Unsigned} + n = sizes[1] % U + q = Base.udiv_int(p, n) + return (p - q * n, _cartesian_offsets(q, Base.tail(sizes))...) +end # With `out === nothing`, compute block counts. Otherwise, `block_counts` contains their diff --git a/test/generic/findall.jl b/test/generic/findall.jl index e38a94a5..403c5135 100644 --- a/test/generic/findall.jl +++ b/test/generic/findall.jl @@ -142,6 +142,21 @@ findall_alg(; block_size=nothing, items_per_thread=nothing, max_tasks=nothing, m @test Array(AK.findall(array_from_host(hb); items=v, alg=FINDALL_ALG)) == h[hb] @test Array(AK.findall(array_from_host(hb); items=CartesianIndices(v), alg=FINDALL_ALG)) == findall(hb) + # Cartesian indices with offset and stepped axes + ci = CartesianIndices((2:38, 0:2:56)) + @test Array(AK.findall(array_from_host(hb); items=ci, alg=FINDALL_ALG)) == + ci[findall(vec(hb))] + # ... and the conversion itself, for other axes and on both sides of 32-bit lengths (which + # only 64-bit hosts can index) + for ci in (CartesianIndices((Base.IdentityUnitRange(-2:4), 5:-1:1, 3)), + CartesianIndices((7,)), CartesianIndices(()), + (Sys.WORD_SIZE == 64 ? (CartesianIndices((65_536, 65_535)), + CartesianIndices((65_536, 65_537, 2))) : ())...) + n = length(ci) + for p in unique(clamp.((1, 2, 7, n รท 3, n - 1, n), 1, n)) + @test AK.findall_index(ci, p) == ci[p] + end + end # ... of another shape or on another array, paired by position w = array_from_host(collect(Int32, 1:length(h))) @test Array(AK.findall(pred, v; items=w, alg=FINDALL_ALG)) == findall(sel)