Repository navigation
Use GPUArrays 12 for sorting, scans, reductions and findall - #3356
Merged
Merged
Conversation
maleadt
marked this pull request as ready for review
October 7, 2026 06:17
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3356 +/- ##
==========================================
- Coverage 87.38% 84.94% -2.44%
==========================================
Files 195 190 -5
Lines 20044 19147 -897
==========================================
- Hits 17515 16264 -1251
- Misses 2529 2883 +354 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
maleadt
force-pushed
the
tb/gpuarrays-12
branch
from
October 7, 2026 14:33
b92c7be to
3f1c0b7
Compare
Contributor
|
I don't have the time to look into this, but I'll mention that this seems to cause a large regression in Oceananigans, here's a reproducer without Oceananigans itself # Full reduction of a lazy, device-adaptable `AbstractArray` into a one-element `CuArray` with
# `Base.mapreducedim!`, the pattern Oceananigans uses for dot products and norms of fields: the
# source computes `a[i, j, k] * b[i, j, k] * mask[i, j, k]` on the fly from halo-padded storage.
#
# usage: julia --project=<env> lazy_reduction.jl [N]
using CUDA, Adapt, Printf
using GPUArrays: GPUArrays
const N = length(ARGS) ≥ 1 ? parse(Int, ARGS[1]) : 64
const H = 4 # halo width
# Lazy elementwise product of two padded arrays, masked: interior index (i, j, k) reads
# the padded storage at (i + H, j + H, k + H). Neither dense nor strided, like a lazy field operation.
struct MaskedProduct{T, A, M} <: AbstractArray{T, 3}
a :: A
b :: A
mask :: M
size :: NTuple{3, Int}
end
MaskedProduct(a, b, mask, size) = MaskedProduct{eltype(a), typeof(a), typeof(mask)}(a, b, mask, size)
Base.size(p::MaskedProduct) = p.size
Base.@propagate_inbounds Base.getindex(p::MaskedProduct, i::Int, j::Int, k::Int) =
p.a[i + H, j + H, k + H] * p.b[i + H, j + H, k + H] * p.mask[i, j, k]
Adapt.adapt_structure(to, p::MaskedProduct) =
MaskedProduct(adapt(to, p.a), adapt(to, p.b), adapt(to, p.mask), p.size)
function setup(T=Float64)
a = CuArray(rand(T, N + 2H, N + 2H, N + 2H))
b = CuArray(rand(T, N + 2H, N + 2H, N + 2H))
mask = CuArray(T[k > N ÷ 4 for i in 1:N, j in 1:N, k in 1:N]) # bottom quarter "immersed"
lazy = MaskedProduct(a, b, mask, (N, N, N))
dense = a[H+1:H+N, H+1:H+N, H+1:H+N] .* b[H+1:H+N, H+1:H+N, H+1:H+N] .* mask
reference = sum(Array(dense))
return (; lazy, dense, reference)
end
variants(s) = (
"dense CuArray (control)" => s.dense,
"lazy AbstractArray" => s.lazy,
"Broadcasted of the lazy" => Broadcast.instantiate(Broadcast.broadcasted(identity, s.lazy)),
)
function reduce_into!(r, src)
fill!(r, 0)
Base.mapreducedim!(identity, +, r, src)
return r
end
function measure(name, src, reference; calls=100)
r = CUDA.zeros(Float64, 1, 1, 1)
reduce_into!(r, src); CUDA.synchronize() # compile
ok = isapprox(Array(r)[1], reference; rtol=1e-12)
bytes = @allocated (reduce_into!(r, src); CUDA.synchronize())
host = @elapsed (for _ in 1:calls; reduce_into!(r, src); end) # host time to enqueue
CUDA.synchronize()
wall = @elapsed (for _ in 1:calls; reduce_into!(r, src); end; CUDA.synchronize())
prof = CUDA.@profile trace=true reduce_into!(r, src)
names = prof.host.name
launches = count(n -> occursin("aunchKernel", n), names)
allocs = count(n -> occursin("MemAlloc", n), names)
syncs = count(n -> occursin("ynchronize", n), names) # includes the profiler's own
gpu = sum(prof.device.stop .- prof.device.start) * 1e6
kernels = unique(first.(prof.device.name, 40))
@printf("%-25s correct=%-5s wall=%8.1f µs/call host=%7.1f µs/call GPU=%8.1f µs host alloc=%6d B launches=%d device allocs=%d syncs=%d\n",
name, ok, wall / calls * 1e6, host / calls * 1e6, gpu, bytes, launches, allocs, syncs)
println(" kernels: ", join(kernels, " | "))
end
function main()
println("CUDA.jl ", pkgversion(CUDA), ", GPUArrays ", pkgversion(GPUArrays), ", ", CUDA.name(CUDA.device()), ", N = $N")
s = setup()
for (name, src) in variants(s)
measure(name, src, s.reference)
end
end
main()We may be able to change the source code not to depend on this, but I wanted to flag that there are workloads which may be negatively affected |
Contributor
CUDA.jl BenchmarksDetails
This comment was automatically generated by workflow using github-action-benchmark. |
Member
Author
maleadt
force-pushed
the
tb/gpuarrays-12
branch
from
October 8, 2026 20:01
3f1c0b7 to
374a459
Compare
GPUArrays 12 implements these once for every back-end on AcceleratedKernels, following Base's semantics, and no longer calls the GPUArrays.mapreducedim! hook. Remove CUDA's own implementations, which would otherwise shadow the generic ones (and in the case of findall, be ambiguous with them). The warp- and block-level reduction helpers stay, as the library packages use them. This drops CUDA.QuickSort and CUDA.BitonicSort and the internal scan!, which is technically breaking but has no known users.
GPUArrays 12.1 defines `Base.dataids` and `Base.mightalias` for every GPU array from where its elements live. Implement that hook instead of our own methods. Empty arrays now alias nothing (the testsuite's empty-view check failed), and Base's check for two SubArrays no longer converts their parents to pointers, which took stream ownership of their memory.
maleadt
force-pushed
the
tb/gpuarrays-12
branch
from
October 9, 2026 11:05
374a459 to
a4872f0
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ports CUDA.jl to GPUArrays 12. GPUArrays 12 implements Base's sorting, scans, reductions,
findall, logical indexing andreverseonce for every GPU array, on AcceleratedKernels 0.5 (JuliaGPU/GPUArrays.jl#790), and no longer calls theGPUArrays.mapreducedim!hook. This deletes CUDA.jl's own implementations, about 2,000 lines:indexing.jl,accumulate.jl(including the internalscan!),reverse.jlandsorting.jl(withCUDA.QuickSortandCUDA.BitonicSort) are removed fromCUDACore/src.findall(f, ::AnyCuArray)was also ambiguous with GPUArrays'findall(::Fix2{typeof(in)}, …).mapreduce.jlkeeps onlyreduce_warp/reduce_block, which cuBLAS and cuSPARSE use, asreduction.jl.Tested on an RTX 5080 (Julia 1.13): all of
core, cuBLAS, cuSPARSE and the GPUArrays testsuite pass, exceptcholesky(::Diagonal)(3 errors). That's a closure capturing aComposedFunctionthat Adapt doesn't adapt inferably, fixed in JuliaGPU/Adapt.jl#124.Timings, CUDA 6.4.2 vs this branch, in ms, best of 20 (sorts best of 10, including a copy):
sum, 1e7 Float32sum, 1000×10000,dims=1sum, 1000×10000,dims=2sum, 1000×10,dims=2sort!, 1e7 Int32sort!, 1e7 Float32sortperm, 1e7 Float32cumsum, 1e7 Float32findall(>(0.5f0), x), 1e7x[x .> 0.5f0], 1e7Small reductions along a dimension are slower; that's tracked in JuliaGPU/AcceleratedKernels.jl#135.