Repository navigation
Reduce whole arrays without a view of their source - #190
Merged
Merged
Conversation
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
force-pushed
the
ab/reduce-without-owning-view
branch
from
October 28, 2026 09:41
3f55e88 to
3dd5372
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
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.
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 aCuArray, GPUArrays builds that view withderive, which makes a newCuArrayholding a reference to the source's memory until the finalizer runs. Since GPUArrays 12 sends Base reductions through AK, this affects plainsum,maximum,countand friends, and it applies tocommutative=falseanddimsreductions into a single output too: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 alongdimsinto 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 aCuVectortakes ~250 ms instead of ~100 ms. The generated code is unchanged apart from computing the length from both dimensions, andsum/maximumof aCuVectorand aCuMatrix(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 checksused_memory()afterunsafe_free!for a vector, a matrix withdims=(1, 2),commutative=falseand 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.