Repository navigation
Reduce lazy sources along dims with the launch shapes of dense ones - #154
Merged
Merged
Conversation
A source without a dense buffer, such as a lazy array or a Broadcasted object, was reduced along dims by a kernel with one thread per output, whatever the number of outputs. Reducing such a source into one element therefore ran on a single thread: a dot product of two lazy 128^3 fields, as Oceananigans computes them, took 220 ms instead of 23 us. Index these sources linearly, as a dense array of their size, so that they take the launch shapes of a dense array: the whole-array kernels for one output, and the cooperative kernels for few outputs. Only the shape with one output per thread keeps the Cartesian-indexed kernel, which needs no index decoding. A lazy array is indexed through a reshape, which decodes linear indices with multiplicative inverses, and is never read with the by-block kernel's vector loads. This also fixes whole-array reductions (`dims=:`) of vectors whose indices do not start at 1, which viewed their elements from index 1.
A reduction along dims into one element calls the whole-array kernels, whose host code was inferred for every new `f`, `op` and element type, whether the reduction took that path or not. Call it through `Base.inferencebarrier`, like the other launches of `mapreduce_nd!`. This saves about 0.05 s of the first call of a reduction along dims with a new `f` on CUDA, and makes up for the inference that lazy sources now share with dense ones.
A two-input reduction reads a fused Broadcasted source, which has no dense buffer. Measure it into one output and into few outputs of many elements each, the cases that used to run one thread per output.
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
1 participant
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.
mapreduce_nd!reduced any source without a dense buffer (a custom lazyAbstractArray, or aBroadcastedon Julia 1.12+) with one thread per output, whatever the number of outputs. A reduction into one element therefore ran on a single GPU thread. Since GPUArrays 12 runs Base's reductions on AK, this hits Oceananigans' dot products and norms, which useBase.mapreducedim!into a 1×1×1 array with a lazy source (reported on JuliaGPU/CUDA.jl#3356, cc @giordano).Such sources are now indexed linearly, as a dense array of their size, so they take the same launch shapes as a dense array: the whole-array kernels for one output, and the cooperative kernels (multigroup, by-block, columns, tiled) for few outputs. Only the one-output-per-thread shape keeps the Cartesian-indexed kernel, because decoding linear indices made it 10–60% slower there. Lazy arrays are indexed through
vec, a reshape that decodes indices with multiplicative inverses, and are never read with the by-block kernel's vector loads.A second commit calls the whole-array launch through
Base.inferencebarrier, as #153 did for the other launches. A third adds a benchmark of a fused two-input source.GPU time on an RTX 5080 (CUDA.jl #3356, Julia 1.13), for a lazy N³ source
a[i+H,j+H,k+H] * b[...] * mask[i,j,k]:dims=(1,2), N = 128dims=(2,3), N = 128dims=(1,3), N = 128dims=1/2/3, N = 128Broadcastedsources alongdimsimprove the same way:mapreduce(*, +, a, b; dims=(1,2))on two 128³ arrays takes 56 µs instead of 1.9 ms. The first call of a reduction with a newftakes about 0.06 s less for dense sources and up to 0.08 s more for lazy ones, which now infer the same paths as dense sources. A reduction into one element allocates 176 more bytes on the host.This also fixes whole-array reductions of vectors whose indices don't start at 1.