Skip to content

Finish whole-array reductions on the device - #139

Merged
maleadt merged 1 commit into
tb/reduce-smallfrom
tb/reduce-device-finish
Oct 2, 2026
Merged

maleadt merged 1 commit into
tb/reduce-smallfrom
tb/reduce-device-finish

Conversation

@maleadt

@maleadt maleadt commented Oct 1, 2026 •

Copy link
Copy Markdown
Member

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 as mapreducedim! into a one-element array, sum!(r, A) or sum(A; dims=(1, 2)), used the general dims machinery instead, and that was up to 20% slower. GPUArrays needs both: a scalar for sum(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_thread elements (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-block dims shape. That shape caps itself at 256 blocks and has no unrolled loads, so a reduction of 10^8 Float32s 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 DeviceReduce works:

  1. A first pass of at most target_blocks blocks (256 by default, a per-backend tuning value). Each block strides over tiles of the input and produces one partial result.
  2. A second pass of a single block, which reduces those partial results.

A small input that fits in one block takes one launch. Every pass stores its values through the same finishing step the dims kernels use. So the last pass can apply init, or fold into R'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! and mapreduce(...; dims) use this path when there is a single output and the source is a Broadcasted or 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_thread now defaults to 64 bytes of accumulator per thread: 16 for 4-byte types, between 4 and 16 in general. A backend tuning that sets no items_per_thread gets this default. On the 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.

Numbers

RTX 5080, Float32 sums, 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. The mapreducedim! column includes copying the result to the host, so it compares directly with AK.mapreduce:

n AK.mapreduce, main this PR mapreducedim! into 1 element, main this PR CUDA.jl mapreducedim! + copy
10^4 14.5 14.8 29.1 15.4 15.6
10^6 18.4 12.1 28.7 15.6 16.8
10^8 459.7 454.6 550.7 458.1 470.0

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 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:

n M1 main M1 this PR Iris Xe main Iris Xe this PR
10^4 367 / 398 372 / 373 338 / 321 382 / 381
10^6 502 / 484 397 / 396 520 / 499 463 / 479
10^8 7197 / 6856 7081 / 7101 19724 / 20280 20559 / 20374

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. init is still applied once, partial results keep the accumulator type until the final store, and R is never used for intermediate values. The visible differences are implementation details:

  • workspace_size returns smaller partial-result buffers.
  • Auto resolves items_per_thread differently.
  • target_blocks now also caps the first pass of a whole-array reduction.
  • Rounding can differ from earlier versions, because the reduction tree changed.

The new tests reduce into a single output across one tile, one pass and two passes. They cover init, folding, overwrite, an accumulator wider than R whose partial results do not fit in R, 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).

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 |
@maleadt
maleadt changed the base branch from main to tb/reduce-small October 1, 2026 14:26
@maleadt
maleadt added this pull request to stack #141 October 1, 2026 14:26
@maleadt
maleadt merged commit 253bd4e into main Oct 2, 2026
37 of 39 checks passed
@maleadt
maleadt deleted the tb/reduce-device-finish branch October 2, 2026 09:01
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.

1 participant