Repository navigation
Finish whole-array reductions on the device - #139
Merged
Merged
Conversation
The whole-array kernel launched one block per tile of `block_size * items_per_thread` elements, then reduced the partial results in further passes until one was left, and returned it with a scalar copy. That took three launches at 10^6 elements, and its result could only reach the host, so a reduction whose result must stay on the device, `mapreducedim!` into a one-element array or `sum(A; dims=(1, 2))`, used the `multigroup` shape instead, which caps its blocks at 256 without unrolled loads: at 10^8 elements 20% slower. Reduce in at most two launches, as CUB's `DeviceReduce` does: a first pass of at most `target_blocks` blocks (256 by default), each striding over tiles of the input, and a second pass of one block over their partial results. Each pass stores its values through `_finish`, so the last one can apply `init` (or fold into `R`) and write the result into a device array; for a host result it stores the bare partial result, which is copied back and finished on the host as before. Reductions along `dims` into a single output of a contiguous (or `Broadcasted`) source now use these kernels. With few blocks, the first pass needs more loads in flight per thread, so a tuning that leaves `items_per_thread` unset (now the default) loads 64 bytes of accumulators per thread and tile (16 for 4-byte types, from 4 to 16); on an RTX 5080 16 is within 1% of the best setting for `Float32` from 10^4 to 10^8 elements. The workspace of a whole-array reduction shrinks from twice the number of tiles to `target_blocks + 1` partial results. `Float32` sums, minimum of 50 synchronized runs (`CUDA.@elapsed`), RTX 5080, µs; `mapreducedim!` into a one-element array includes the copy of the result to the host: | n | `AK.mapreduce` before | after | `mapreducedim!` before | after | |-------|-----------------------|-------|------------------------|-------| | 10^4 | 14.5 | 14.8 | 29.1 | 15.4 | | 10^6 | 18.4 | 12.1 | 28.7 | 15.6 | | 10^8 | 459.7 | 454.6 | 550.7 | 458.1 |
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.
Part of #135. Stacked on #138, which has to be merged first; this PR only adds the device-side finish.
AK has two ways to reduce a whole array. The whole-array kernel behind
AK.mapreduce(f, op, A)returns a host value. A reduction whose result has to stay on the device, such asmapreducedim!into a one-element array,sum!(r, A)orsum(A; dims=(1, 2)), used the generaldimsmachinery instead, and that was up to 20% slower. GPUArrays needs both: a scalar forsum(A)and a device array for the others.What was wrong
The whole-array kernel launched one block per tile of
block_size * items_per_threadelements (512 by default). It then reduced the partial results in further passes until one value was left, and copied that value to the host. At 10^6 elements that took three launches plus the copy. It could not leave its result in a device array.mapreducedim!into a single output fell through to the multi-blockdimsshape. That shape caps itself at 256 blocks and has no unrolled loads, so a reduction of 10^8Float32s took 550 µs instead of the whole-array kernel's 458 µs.What changes
Whole-array reductions now take at most two launches, the way CUB's
DeviceReduceworks:target_blocksblocks (256 by default, a per-backend tuning value). Each block strides over tiles of the input and produces one partial result.A small input that fits in one block takes one launch. Every pass stores its values through the same finishing step the
dimskernels use. So the last pass can applyinit, or fold intoR's previous value, and write the result into a device array. When the caller wants a host value, the last pass stores the bare partial result, and it is copied back and finished on the host as before.mapreducedim!andmapreduce(...; dims)use this path when there is a single output and the source is aBroadcastedor its elements form one contiguous range of memory.Because the first pass now launches few blocks, each thread needs more loads in flight to reach full bandwidth.
items_per_threadnow defaults to 64 bytes of accumulator per thread: 16 for 4-byte types, between 4 and 16 in general. A backend tuning that sets noitems_per_threadgets this default. On the RTX 5080, 16 is within 1% of the best setting forFloat32from 10^4 to 10^8 elements. The workspace of a whole-array reduction shrinks from twice the number of tiles totarget_blocks + 1partial results.Numbers
RTX 5080,
Float32sums, Julia 1.13. Minimum of 50 synchronized runs (CUDA.@elapsed), variants interleaved, best of 3 rounds, on a GPU shared with at most one other process (up to 14% utilization). µs. Themapreducedim!column includes copying the result to the host, so it compares directly withAK.mapreduce:AK.mapreduce, mainmapreducedim!into 1 element, mainmapreducedim!+ copyOn Metal and oneAPI, take the numbers below as indicative only. Both machines were busy with other jobs during every run: another session's Metal tests on the M1, and CI jobs on the Iris Xe. Wall clock, minimum of 30 runs, best of 2 to 4 rounds, µs, for
AK.mapreduce/mapreducedim!+ copy:On these machines the differences at 10^4 and 10^8 are within the noise: a build that runs main's reduction code measured 356 / 398 µs on the Iris Xe at 10^4 in the same session.
Compatibility
There are no API or contract changes, and this can ship in a 0.5.x release.
initis still applied once, partial results keep the accumulator type until the final store, andRis never used for intermediate values. The visible differences are implementation details:workspace_sizereturns smaller partial-result buffers.Autoresolvesitems_per_threaddifferently.target_blocksnow also caps the first pass of a whole-array reduction.The new tests reduce into a single output across one tile, one pass and two passes. They cover
init, folding,overwrite, an accumulator wider thanRwhose partial results do not fit inR, workspace reuse, offset views and fused sources. Tested on CUDA, OpenCL (PoCL), the CPU threads back-end, Metal, oneAPI, and AMDGPU (an APU, where the failures that main already has remain).