Repository navigation
Speed up small reductions and reductions of wide matrices along dims - #138
Merged
Merged
Conversation
Julia does not specialize a method on a function argument that it only
passes on to other calls. The reductions' entry points passed `f` and
`op` straight through, so the host code below them ran unspecialized:
`Base.promote_op` ran inference at run time on every call to find the
accumulator type, and every call after it was a dynamic dispatch. The
dimension segments added more, as vectors of tuples of unknown length.
Annotate the entry points with `where {F, OP}`, keep the segments in
fixed-length padded tuples (dispatching the common layouts, one reduced
and at most one kept segment, statically), and put a function barrier
after the computation of the reduced dimensions in `mapreducedim!`.
Host time of `mapreducedim!(identity, +, R, A; init=0f0)` for a 1000×10
`Float32` matrix (`dims=2`), RTX 5080: 13.9 µs before, 3.7 µs after;
CUDA.jl's own kernel takes 2.6 µs, of which the launch is about 2 µs.
That was the whole gap for small reductions (wall-clock minimum of
synchronized runs, µs):
| shape | dims | CUDA.jl | before | after |
|-----------|------|---------|--------|-------|
| 1000×10 | 1 | 11.8 | 20.0 | 8.1 |
| 1000×10 | 2 | 6.8 | 19.9 | 7.5 |
| 1000×1000 | 1 | 17.8 | 23.1 | 10.6 |
| 1000×1000 | 2 | 23.8 | 21.7 | 9.6 |
A reduction into contiguous outputs along a longer strided dimension, such as `sum(A; dims=2)` of a wide column-major matrix, took the `by_block` shape: one block per output, its threads reading elements a row apart, so no two threads of a warp touched the same cache line. At 1000×100000 it ran at a quarter of the bandwidth, slower than CUDA.jl's kernel; with few outputs it took `multigroup`, as poorly coalesced. The column kernel gives each block up to 32 consecutive outputs (8 to 31 outputs take the next power of two) with the block's other threads as lanes along the reduction, so a warp's loads of one step are consecutive elements; each thread keeps 8 loads in flight, and when the outputs alone do not fill the device, each output's reduction is split across blocks and combined in a second pass, like `multigroup`. It runs blocks of at most 256 threads, as Metal allows its pipeline fewer than 1024. It is used from `columns_min_elements` elements, a new internal tuning value: 2^18 by default, 2^22 on oneAPI, where a second launch costs tens of microseconds; and from 8 outputs, as with 4 it lost to `multigroup` on Metal and oneAPI. `sum(A; dims=2)` of `Float32`, minimum of 50 synchronized runs (`CUDA.@elapsed`), RTX 5080, µs: | shape | CUDA.jl | before | after | |--------------|---------|--------|-------| | 300×10000 | 45.3 | 67.0 | 10.6 | | 1000×10000 | 141.0 | 171.2 | 18.3 | | 1000×100000 | 1459.1 | 1815.2 | 454.6 | | 5000×20000 | 1466.2 | 1849.9 | 460.6 | | 100×1000000 | 1705.9 | 1657.2 | 460.1 | | 32×3000000 | 1452.3 | 4374.9 | 424.7 | | 16×6000000 | 1078.0 | 761.5 | 428.4 | | 8×12000000 | 701.8 | 760.5 | 430.0 |
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.
Part of #135.
GPUArrays is about to send every Base reduction of GPU arrays through AK (JuliaGPU/GPUArrays.jl#790), so AK's reductions have to be at least as fast as the back-ends' own kernels they replace. For small reductions along
dimsthey were not:sum(A; dims=2)of a 1000×10 matrix took 17.9 µs against 4.5 µs for CUDA.jl'smapreducedim!. Reductions of wide matrices alongdims=2were also slower than CUDA.jl's. This PR fixes both.Small reductions: host overhead
The kernels were not the problem. Almost all the time went into host code that ran unspecialized. Julia does not specialize a method on a function argument that it only passes on to other calls. AK's entry points passed
fandopstraight through, so everything below them was compiled for::Function.Base.promote_opthen ran type inference at run time, on every call, to find the accumulator type, and each call after it was a dynamic dispatch. The reduced and kept dimensions added more, stored as vectors of tuples of unknown length.The first commit annotates the entry points with
where {F, OP}and keeps the dimension segments in fixed-length tuples. It also adds a function barrier wheremapreducedim!computes the reduced dimensions, whose number is only known at run time. Host time ofmapreducedim!(identity, +, R, A; init=0f0)on a 1000×10Float32matrix drops from 13.9 µs to 3.7 µs. CUDA.jl's call takes 2.6 µs, about 2 µs of which is the kernel launch.Wide matrices: a column kernel
A reduction into contiguous outputs along a longer, strided dimension, the row sums of a wide column-major matrix for example, used one block per output. Its threads read elements a whole column apart, so the loads of a warp hit as many cache lines as it has threads. At 1000×100000 it ran at a quarter of the memory bandwidth. With fewer than 32 outputs it used the multi-block shape, which reads the input the same way.
The second commit adds a kernel for this layout. Each block takes up to 32 consecutive outputs (8 to 31 outputs take the next power of two) and spreads its other threads along the reduction, so that a warp's loads at each step are consecutive elements. Each thread keeps 8 loads in flight. When there are too few outputs to fill the device, the reduction of each output is split across blocks, and a second pass combines the partial results, as the multi-block shape already does. The kernel runs at most 256 threads per block because Metal does not accept 1024 threads for it. It is chosen from
columns_min_elementselements on, a new internal tuning value: 2^18 by default, and 2^22 on oneAPI, where the second launch costs tens of microseconds. It also needs at least 8 outputs. With 4, it was faster on CUDA but slower than the multi-block shape on Metal and oneAPI.Numbers
RTX 5080,
Float32,init=0f0, Julia 1.13, CUDA.jl 6.4.1. Minimum of 50 synchronized runs per variant, variants interleaved, best of 3 rounds (2 for wall clock). The GPU was shared with at most one other process, at up to 14% utilization. µs.mapreducedim!against CUDA.jl'sGPUArrays.mapreducedim!. For the smallest shapes the table shows wall-clock time of the synchronized call. With other work on the GPU,CUDA.@elapsedcan hide part of the host overhead of these sub-10 µs calls:Larger ones,
CUDA.@elapsed:400 MB in 425-460 µs is close to the card's bandwidth. The 4×25000000 case keeps the multi-block shape.
On Metal and oneAPI, take the numbers below as indicative only. Both machines were busy with other jobs during every run: another session's Metal test suites on the M1, and a stream of CI jobs on the Iris Xe. Wall clock, minimum of 30 synchronized runs, best of 2 to 4 rounds, µs:
On the Iris Xe, 300×10000 is below oneAPI's threshold and runs the same kernel as on main, so its difference is noise.
Compatibility
There are no API or contract changes. The new
ReduceTuningfield is internal, and the reduction tree changes only for the shapes that now take the column kernel. That can change floating-point rounding, which the reduction contract allows.Tested on CUDA, OpenCL (PoCL), the CPU threads back-end, Metal, oneAPI, and AMDGPU (an APU, where the failures that main already has remain). The new tests cover the column kernel in one and two passes, without a neutral element, with
init, folding intoR, anRnarrower than the accumulator, a reused workspace, and a strided view at an offset.