[CUDA] Parallelize ArgMax/ArgMin over wide last axes (#32092)
### Description
`arg_min_max_last_axis()` in `reduction_functions.cu` assigns **one
thread per row** and scans the row serially. That is a reasonable design
for many narrow rows, but it degenerates badly when a graph reduces a
small number of very long rows: the reduction length becomes a serial
dependency chain in a single thread, so the cost grows linearly with the
reduced axis and the rest of the GPU stays idle.
A `[1, 202048]` `ArgMax` (last-axis classification / sampling over a
large vocabulary) is therefore computed by exactly **one** CUDA thread
and takes **~5 ms** on an H200. The same shape costs microseconds once
it is parallelized. This affects any model that reduces a wide last axis
with few rows, and it also leaves a lot on the table for the medium
cases (thousands of rows of a few hundred to a few thousand columns).
This PR adds a cooperative reduction path and keeps the existing serial
kernel for the shapes where it is genuinely better.
### Approach
- A new kernel spreads one row over many warps and, when the row is long
enough, over many blocks:
- thread level scan (grid-strided, 4 elements per thread, coalesced),
- warp level merge with shuffles,
- block level merge through shared memory,
- grid level merge for multi-block rows through a small global buffer,
finalized by the last block arriving for that row.
- The multi-block step reuses the `__threadfence()` + per-row
done-counter structure that `reduce_matrix_columns()` / `reduce_all()`
already use in the same file, so the intermediate buffer follows the
existing sizing and alignment conventions and is allocated by the caller
with `AllocateScratchBuffer` (same as the
`ApplicableMatrixReduction::Columns` path).
- Launch geometry aims for a single wave of threads over the device
while keeping at least four elements per thread. This avoids both
failure modes: an idle device for few wide rows, and oversubscription
(which measures 2-3x slower) for many rows.
- Rows narrower than 128 columns keep the existing one-thread-per-row
kernel; below that width a row cannot fill a warp and the serial kernel
wins. The threshold was picked from a measured crossover sweep.
- No model-specific constants: the dispatch depends only on the matrix
shape and the device's thread capacity.
### Semantics
The new kernel is **bit-identical to the previous kernel for every
input**, including:
- **Ties**: the lowest index wins, matching `select_last_index == 0`.
The CUDA EP already falls back to CPU for `select_last_index == 1`
(`ArgMaxOrArgMinNeedFallbackToCPU`), so this is the only case to
support.
- **NaN**: NaN never wins a comparison, and the previous kernel seeded
its accumulator with element 0, so a leading NaN yields index 0 for the
whole row. That behavior is preserved explicitly.
- **Infinities**: the reduction identity is `-inf` / `+inf`, not the
lowest/highest finite value. This matters: with a finite identity, a row
of `[-inf, ..., lowest_finite, ...]` would drop the real maximum. Rows
that are entirely `-inf`, `+inf` or NaN return index 0, as before.
- `axis`, `keepdims`, negative axes, opset variants and the int64 output
are untouched, they are handled above this function.
- All registered input types are covered (`MLFloat16`, `float`, `double`
for ArgMax/ArgMin) and the fallbacks for non-last-axis reductions,
`n`/`m` beyond `int` range and empty reductions are unchanged.
### Results
H200 (sm90), CUDA 13.3, fp32, per call including the launch, measured
with the benchmark added in this PR
(`ReductionFunctionsTest.DISABLED_ArgMinMaxLastAxisPerf`), "before"
measured with the same harness on the serial kernel:
| rows | cols | before | after | speedup | scratch |
|---:|---:|---:|---:|---:|---:|
| 1 | 32000 | 695.6 us | 5.6 us | 125x | 267 B |
| 1 | 50257 | 1201.1 us | 5.7 us | 210x | 411 B |
| 1 | 128256 | 3009.7 us | 5.9 us | 511x | 1019 B |
| 1 | 151936 | 3833.8 us | 6.2 us | 620x | 1203 B |
| 1 | 200000 | 5037.1 us | 6.1 us | 821x | 1579 B |
| 1 | 202048 | 5056.5 us | 6.2 us | 818x | 1595 B |
| 1 | 262144 | 6335.5 us | 6.2 us | 1018x | 2067 B |
| 2 | 202048 | 4853.4 us | 6.5 us | 746x | 3175 B |
| 4 | 202048 | 4759.6 us | 7.5 us | 635x | 6335 B |
| 8 | 202048 | 5118.4 us | 8.5 us | 603x | 8495 B |
| 16 | 202048 | 5805.6 us | 9.1 us | 640x | 8527 B |
| 32 | 202048 | 7357.7 us | 10.8 us | 681x | 8591 B |
| 4096 | 128 | 18.8 us | 5.0 us | 3.75x | 0 |
| 4096 | 1024 | 135.6 us | 6.5 us | 21x | 0 |
| 4096 | 4096 | 552.2 us | 20.8 us | 27x | 0 |
| 65536 | 256 | 70.7 us | 28.3 us | 2.50x | 0 |
| 65536 | 1024 | 270.1 us | 67.7 us | 3.99x | 0 |
| 4096 | 8 | 2.3 us | 2.3 us | 1.01x | 0 |
| 4096 | 32 | 6.2 us | 6.2 us | 0.99x | 0 |
| 4096 | 64 | 10.5 us | 10.6 us | 0.99x | 0 |
| 65536 | 64 | 18.9 us | 19.0 us | 1.00x | 0 |
Narrow rows are unchanged because they keep the existing kernel. `half`
and `double` behave the same way (`[1, 202048]`: 2432 us -> 6.1 us for
`half`, 6405 us -> 6.5 us for `double`).
Profile of `[1, 202048]` fp32 (Nsight Systems, occupancy from
`cudaOccupancyMaxActiveBlocksPerMultiprocessor`):
| | before | after |
|---|---|---|
| kernel launches | 1 | 1 (+ one 4 byte memset) |
| GPU time | ~5056 us | 3.8 us kernel + 0.8 us memset |
| grid / block | 1 x 256 (1 active thread) | 197 x (32, 8) |
| registers / thread | - | 32 |
| theoretical occupancy | - | 100% |
| scratch | 0 | 1595 B |
For reference on the same GPU and shape, `cub::DeviceReduce::ArgMax`
needs 2 kernels, 2.74 us of GPU time and 42495 bytes of temporary
storage. The kernel here is a segmented (per row) reduction that also
has to serve the many-rows shapes, uses ~26x less scratch and one launch
instead of two. Scratch is bounded by the device thread capacity (a
multi-block launch only happens when `rows < capacity / 256`), so it
stays in the low kilobytes for any shape.
### Testing
-
`onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc`:
new `ArgMinMaxLastAxis*` cases covering every dispatch path (narrow /
single block / multi block / more rows than grid rows),
`half`/`float`/`double`, ArgMax and ArgMin, ties, NaN (leading, middle,
all), `+/-inf`, all-equal rows, finite extremes, unaligned and
undersized intermediate buffers, a dirtied intermediate buffer, and CUDA
graph capture + repeated replay.
- `onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc`:
operator level ArgMax/ArgMin tests with a wide last axis (`{3, 40000}`
and `{2, 202048}`, with `keepdims` on and off and a negative axis) so
the CUDA EP scratch allocation path is covered end to end.
- `onnxruntime_provider_test --gtest_filter='*ArgMax*:*ArgMin*'`: 32/32
pass.
- `onnxruntime_provider_test --gtest_filter='*Reduce*'`: 376/376 pass.
- `CUDA_EP_Unittest.All` (CUDA EP internal tests): 76/76 pass.
- `compute-sanitizer` (`memcheck`, `racecheck`, `synccheck`) clean on
the multi-block and row-loop paths.
- Differential test against the previous kernel's exact sequential
semantics over a dense width sweep around the dispatch threshold and
around power-of-two boundaries, several row counts, all three dtypes,
random data and a `{+/-inf, NaN, +/-0.0, ...}` alphabet, repeated with
several simulated device capacities: no mismatches.
### Review follow-up
- The thread level scan originally advanced `id` in `int`. Since a row
is admitted up to `INT_MAX` columns and the step reaches `4 * 65536 =
262144`, that overflowed to negative positions for the widest rows,
reading out of bounds and never terminating. Fixed by expressing the
scan with differences (`remaining = num_cols - id`, `delta = i *
threads`, positions only formed after `delta < remaining`), which keeps
every value inside `[0, num_cols)`. int64 positions were measured as the
alternative and were 4-38% slower on the targeted shapes, so the
difference based form was kept; no shape is dropped either way. New
tests: `[1, INT_MAX]` fp16 (memory gated), widths straddling the widest
scan step, and buffer sizing at `INT_MAX`.
### Possible follow-up
The thread level scan uses scalar loads. Vectorized loads would shave
roughly another microsecond off the very wide single row case; it was
left out here to keep the alignment handling out of a generic kernel
that has to serve `half`, `float` and `double`.
---------
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
onnxruntime
d4d792f5 - [CUDA] Parallelize ArgMax/ArgMin over wide last axes (#32092)
