Repository navigation
Add a commutative keyword to reduce non-commutative operators in order - #165
Merged
Merged
Conversation
AK's GPU reductions combine elements in an order that requires a commutative operator: each thread folds every N-th element, and blocks combine threads with a halving tree. Associative but non-commutative operators, such as matrix products, give wrong results from 3 elements on, while Base's `reduce` keeps the elements in order (#164). `commutative=false` on `mapreduce`, `reduce` and `mapreducedim!` now selects kernels that combine elements in order: every partial result reduces a run of consecutive elements, and only neighbouring partial results are combined, the earlier on the left. A new kernel serves whole arrays and most reductions along `dims`. Each block takes a contiguous range of tiles, its threads reduce runs of consecutive elements (with vector loads where the source is dense), and the scans' block scan combines the runs. Strided reductions into contiguous outputs use the columns kernel with contiguous ranges per lane, and short slices keep one thread per output. Group counts and scratch are those of the commutative shapes, so workspaces do not depend on the keyword. The default stays `commutative=true`, as in CUB; front ends that implement Base's semantics, such as GPUArrays, choose per operator. On an RTX 5080 with KernelAbstractions 0.10, `commutative=false` takes 1.03x to 1.13x as long as the default for Float32 sums of 10^5 to 10^8 elements, and 0.74x to 1.23x along `dims` of large arrays. Without sub-group support in the block scan (KernelAbstractions 0.9) it takes up to 1.5x as long.
maleadt
force-pushed
the
tb/ordered-reduce
branch
from
October 20, 2026 09:23
503706e to
1db217f
Compare
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.
Fixes #164.
AK's GPU reductions assume a commutative operator: each thread folds every N-th element, and blocks combine threads with a halving tree. An associative but non-commutative
op(matrix products, "first of equal values", affine-map composition) gives wrong results from 3 elements on. CUDA.jl's old kernel hid this below ~100k elements, which is the threshold reported in CUDA.jl#1058.This adds
commutative::Bool=truetomapreduce,reduceandmapreducedim!. Withcommutative=false, every partial result reduces a run of consecutive elements, and only neighbouring partials are combined, the earlier one on the left:dimsgo through a new kernel,_mapreduce_ordered!. Each block takes a balanced, contiguous range of tiles; each thread reduces K consecutive elements (with vector loads for dense sources); the runs are combined with the scans' block scan, which uses sub-groups on KA 0.10;Group counts and scratch are those of the commutative shapes, so workspaces don't depend on the keyword.
The default stays commutative, as in CUB. GPU reduction libraries are commutative-only almost without exception; Futhark's
reducevsreduce_command the open CCCLReduceNonCommutativePR (NVIDIA/cccl#11818) are the exceptions. Front ends that implement Base semantics choose per operator; see JuliaGPU/GPUArrays.jl#822.Performance on an RTX 5080 (KA 0.10,
commutative=falserelative to the default):dimsof large arraysTests: a new "non-commutative reductions" testset uses 2x2
UInt32matrix products and packed affine maps, so the vector-load paths are order-sensitive too. It covers every routing transition, views, lazy sources,init/neutral/overwrite, and workspace reuse. Withcommutative=true, 173 of its 223 tests fail. It passes on CUDA, on the KA CPU backend, and on the host path. It has not been tested on AMDGPU, oneAPI, Metal or OpenCL.