Skip to content

Reduce whole arrays without a view of their source - #190

Merged
maleadt merged 2 commits into
JuliaGPU:mainfrom
AzamatB:ab/reduce-without-owning-view
Oct 29, 2026
Merged

maleadt merged 2 commits into
JuliaGPU:mainfrom
AzamatB:ab/reduce-without-owning-view

Conversation

@AzamatB

@AzamatB AzamatB commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

unsafe_free! of a GPU array doesn't free its memory if the array was reduced first. The whole-array kernels read their source through @view src[begin:end]. For a CuArray, GPUArrays builds that view with derive, which makes a new CuArray holding a reference to the source's memory until the finalizer runs. Since GPUArrays 12 sends Base reductions through AK, this affects plain sum, maximum, count and friends, and it applies to commutative=false and dims reductions into a single output too:

a = CuArray(rand(Float32, 10^6))
sum(a)
CUDA.synchronize(); used = CUDA.used_memory()
CUDA.unsafe_free!(a); CUDA.synchronize()
used - CUDA.used_memory()   # 0 on main, 4000000 with this PR

The kernels index their source linearly from 1, so for a source that already indexes that way (IndexLinear, no offset axes), the view was just another name for the same elements. Such sources are now passed through unchanged. Sources with Cartesian indexing still get the view: it reshapes them, so linear indices are decoded with multiplicative inverses. A reduction along dims into a single output of a wrapper that covers a whole dense array, e.g. sum(transpose(v); dims=2), viewed the wrapped array the same way, and now passes that array instead. Leaving the view in and having CUDA.jl free it early isn't an option, because a derived array is supposed to keep its parent alive. The fix is for AK to not create the view in the first place.

Cost: an N-d array now compiles its own kernel specialization instead of sharing the vector's. On an RTX 5080, the first mapreduce(f, +, ::CuMatrix) after the same call on a CuVector takes ~250 ms instead of ~100 ms. The generated code is unchanged apart from computing the length from both dimensions, and sum/maximum of a CuVector and a CuMatrix (10^4 and 10^8 elements) run at the same speed as before.

Other operations (reductions along dims, scans, sorts, findall, broadcast sources) don't view the user's array and were already fine. The new CUDA test checks used_memory() after unsafe_free! for a vector, a matrix with dims=(1, 2), commutative=false and the transpose case, and fails on main. A wrapper that covers only part of its parent still gets a view, which would take an offset in the kernels to avoid.

AzamatB and others added 2 commits October 28, 2026 10:33
The whole-array kernels read their source through `@view src[begin:end]`.
For a GPU array, GPUArrays makes that view with `derive`: a new array that
shares the source's memory and holds a reference to it until it is
garbage-collected. `unsafe_free!` of an array after any whole-array
reduction (`sum`, `maximum`, `AK.reduce`, ...) therefore no longer freed
its memory.

Pass a source the kernels can index linearly from 1 as it is, and keep the
view (a reshape) only for sources with Cartesian indexing. This costs N-d
arrays a kernel specialization of their own, instead of sharing the
vector's.
A reduction along `dims` into a single output of a wrapper of a dense
array, e.g. `sum(transpose(v); dims=2)` of a vector `v`, passed the
whole-array kernels a view of the wrapped array, which for a GPU array
holds on to its memory as the previous commit describes. Where the
wrapper covers all of that array, pass the array itself.
@maleadt
maleadt force-pushed the ab/reduce-without-owning-view branch from 3f55e88 to 3dd5372 Compare October 28, 2026 09:41
@maleadt
maleadt merged commit ceb8f8f into JuliaGPU:main Oct 29, 2026
17 checks passed
@AzamatB
AzamatB deleted the ab/reduce-without-owning-view branch October 29, 2026 12:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants