From 619f11b9760c4068a2648ad62d4e4ed1ac0ee455 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 1 Oct 2026 15:37:20 +0200 Subject: [PATCH 1/2] Convert linear to Cartesian indices with unsigned divisions in findall MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `findall` with `CartesianIndex` items (`keys(A)`, the default for matrices) converted each selected position with Base's `CartesianIndices` indexing, which divides with checked signed 64-bit divisions. GPUs emulate those in long instruction sequences, and in the scatter pass every thread converts the positions of its run one after the other, so Cartesian `findall` took a fixed 13-19 µs more than with `Int` items, even on small masks. Divide unsigned, in 32 bits where the length fits, and without the zero check: no axis is empty where a position exists. `findall` of a mask, half true, minimum of 50 synchronized runs (`CUDA.@elapsed`), RTX 5080, µs: | mask | `keys` before | after | `LinearIndices` | |-----------|---------------|-------|-----------------| | 100×10 | 38.2 | 24.2 | 23.2 | | 1000×100 | 42.8 | 25.4 | 24.0 | | 4096×4096 | 438.2 | 268.2 | 185.3 | The remaining difference at 4096×4096 is the larger output: 16-byte indices instead of 8-byte ones. --- src/findall.jl | 22 ++++++++++++++++++++++ test/generic/findall.jl | 13 +++++++++++++ 2 files changed, 35 insertions(+) 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..bfa69730 100644 --- a/test/generic/findall.jl +++ b/test/generic/findall.jl @@ -142,6 +142,19 @@ 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 + for ci in (CartesianIndices((Base.IdentityUnitRange(-2:4), 5:-1:1, 3)), + CartesianIndices((7,)), CartesianIndices(()), + 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) From 96e1e34097b953eee2be69ae7ba59d0a6f92ecc4 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 1 Oct 2026 17:16:34 +0200 Subject: [PATCH 2/2] Only test conversions beyond 2^31 elements on 64-bit hosts CartesianIndices of about 2^32 elements have lengths and positions that overflow a 32-bit `Int`, so the reference `ci[p]` in the test threw a BoundsError on x86. --- test/generic/findall.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/generic/findall.jl b/test/generic/findall.jl index bfa69730..403c5135 100644 --- a/test/generic/findall.jl +++ b/test/generic/findall.jl @@ -146,10 +146,12 @@ findall_alg(; block_size=nothing, items_per_thread=nothing, max_tasks=nothing, m 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 + # ... 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(()), - CartesianIndices((65_536, 65_535)), CartesianIndices((65_536, 65_537, 2))) + (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]