Performance of small and whole-array reductions, and of Cartesian `findall`

[maleadt]

GPUArrays is moving every Base reduction of GPU arrays onto AcceleratedKernels (JuliaGPU/GPUArrays.jl#790, on the AK rework #133), and the back-ends' own reduction kernels go. That makes AK's reductions the only ones, and for small reductions they are slower than CUDA.jl's today. This tracks what to optimize.

All times in µs: AK tb/api (#133), CUDA.jl 6.4.0 (main, 4309bcd), an RTX 5080, Julia 1.13; minimum of 30 synchronized runs (CUDA.@elapsed), variants interleaved. The script is below; it needs AK's tb/api branch, and GPUArrays.mapreducedim! on a CuArray is CUDA.jl's kernel.

1. Small reductions are slower than CUDA.jl's. sum of Float32 with init=0f0, CUDA.jl's GPUArrays.mapreducedim! against AK's mapreducedim! (and, for "all", against AK.mapreduce, which is what GPUArrays calls for a scalar result; CUDA.jl's result is copied to the host too):

shape	reduction	CUDA.jl	AK	AK / CUDA.jl
10000	all	16.3	14.9	0.92
1000×10	dims=1	9.5	18.0	1.89
1000×10	dims=2	4.5	17.9	4.00
1000000	all	17.1	18.3	1.07
1000×1000	dims=1	14.9	21.0	1.40
1000×1000	dims=2	20.9	19.1	0.91
100000000	all	469.2	459.8	0.98
1000×100000	dims=1	1410.5	485.2	0.34
1000×100000	dims=2	1490.1	1800.0	1.21

AK loses up to about 4× where there is little work, and wins by about 3× for many short columns (1000×100000 along dims=1). Candidates: the choice of kernel shape and launch configuration for small dims reductions (most shapes run one pass; few outputs with long reductions run two, src/reduce/mapreduce_nd.jl), and fewer launches in the whole-array kernel, which runs several passes and finishes on the host (item 3).

2. mapreducedim! into a one-element array is slower than the whole-array kernel. sum of Float32 with the result on the host:

n	whole-array kernel (AK.mapreduce)	AK.mapreducedim! into a one-element array, then Array(R)
10^4	14.7	27.9
10^6	18.7	29.2
10^8	459.7	547.7

3. The whole-array kernel finishes on the host. mapreduce_1d_gpu (src/reduce/mapreduce_1d_gpu.jl on tb/api) copies the last partial results back and combines them on the host, below switch_below elements or when one is left. So it cannot serve reductions whose result must stay on the device, such as sum(A; dims=(1, 2)) or sum!(r, A), which take mapreducedim! (item 2). A device-side finish, e.g. a last single-block pass writing into a device array, would let every whole-array reduction use the faster kernel.

4. Cartesian findall has a fixed overhead. findall of a mask with half its elements true, selecting CartesianIndexes (items=keys(mask), the default for matrices) against Ints (items=LinearIndices(mask)):

mask	CartesianIndex items (keys)	Int items (LinearIndices)
100×10	38.6	25.1
1000×100	44.6	26.1
4096×4096	439.3	185.0

At 4096² the Cartesian variant costs what Int items plus a separate map to CartesianIndexes cost (185 + about 250 µs, the map measured earlier), so converting the indices in the scatter kernel is no slower than a separate pass. At small sizes, where the conversion costs next to nothing, it still takes 13–19 µs more: a fixed overhead of the Cartesian variant.

Benchmark script

# Numbers for the AK performance issue: CUDA.jl's reduction kernel against AK's, the whole-array
# kernel against `mapreducedim!` into a one-element array, and `findall`'s fixed overhead.
# Minimum of N synchronized runs (CUDA.@elapsed), variants interleaved. RTX 5080.
using CUDA, GPUArrays, Printf
import AcceleratedKernels as AK
const N = 30
function bench(fs::Pair...)
    ts = Dict(k => Inf for (k, _) in fs)
    for (_, f) in fs; f(); end
    CUDA.synchronize()
    for _ in 1:N, (k, f) in fs
        ts[k] = min(ts[k], CUDA.@elapsed f())
    end
    [ts[k] * 1e6 for (k, _) in fs]
end
println("AK ", pkgversion(AK), " (tb/api), CUDA.jl ", pkgversion(CUDA), ", ", CUDA.name(CUDA.device()), ", Julia ", VERSION)

println("\n## sum of Float32, CUDA.jl's `GPUArrays.mapreducedim!` against AK's (μs)\n")
println("| shape | reduction | CUDA.jl | AK | AK / CUDA.jl |")
println("|---|---|---|---|---|")
for n in (10^4, 10^6, 10^8)
    for (shape, rshape, name) in (((n,), (1,), "all"),
                                  ((1000, n ÷ 1000), (1, n ÷ 1000), "`dims=1`"),
                                  ((1000, n ÷ 1000), (1000, 1), "`dims=2`"))
        A = CUDA.rand(Float32, shape...)
        R1 = CUDA.zeros(Float32, rshape...); R2 = CUDA.zeros(Float32, rshape...)
        if name == "all"   # GPUArrays uses the whole-array reduction, whose result is on the host
            tc, ta = bench("c" => () -> (GPUArrays.mapreducedim!(identity, +, R1, A; init=0f0); Array(R1)[1]),
                           "a" => () -> AK.mapreduce(identity, +, A; init=0f0))
        else
            tc, ta = bench("c" => () -> GPUArrays.mapreducedim!(identity, +, R1, A; init=0f0),
                           "a" => () -> AK.mapreducedim!(identity, +, R2, A; init=0f0))
        end
        @printf("| %s | %s | %.1f | %.1f | %.2f |\n", join(shape, "×"), name, tc, ta, ta / tc)
    end
end

println("\n## Whole-array sum of Float32 with the result on the host (μs)\n")
println("| n | whole-array kernel (`AK.mapreduce`) | `AK.mapreducedim!` into a one-element array, then `Array(R)` |")
println("|---|---|---|")
for n in (10^4, 10^6, 10^8)
    v = CUDA.rand(Float32, n); R = CUDA.zeros(Float32, 1)
    tw, tr = bench("w" => () -> AK.mapreduce(identity, +, v),
                   "r" => () -> (AK.mapreducedim!(identity, +, R, v; overwrite=true); Array(R)[1]))
    @printf("| 10^%d | %.1f | %.1f |\n", round(Int, log10(n)), tw, tr)
end

println("\n## `findall` of a mask, half true (μs)\n")
println("| mask | `CartesianIndex` items (`keys`) | `Int` items (`LinearIndices`) |")
println("|---|---|---|")
for (m1, m2) in ((100, 10), (1000, 100), (4096, 4096))
    mask = CuArray(rand(m1, m2) .< 0.5)
    tc, tl = bench("c" => () -> AK.findall(mask), "l" => () -> AK.findall(mask; items=LinearIndices(mask)))
    @printf("| %d×%d | %.1f | %.1f |\n", m1, m2, tc, tl)
end