From fef2239db30cfa90f06bb5c9fa23f2bfdd3ec3db Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 19 Oct 2026 13:46:59 +0200 Subject: [PATCH] Reduce with operators not known to commute in element order AcceleratedKernels' reductions combine elements in any order, so a non-commutative operator, such as one keeping the first of equal values, gave wrong results on large arrays (JuliaGPU/AcceleratedKernels.jl#164, was JuliaGPU/CUDA.jl#1058). Base keeps the elements in order, so pass `commutative=false` to AcceleratedKernels (0.5.8) unless the operator is known to commute on the types involved: Base's arithmetic, bitwise and min/max operators on Base's number types, and this package's own reducers, which become named so that they can be marked. --- Project.toml | 2 +- src/host/indexing.jl | 32 +++++++++++++++-------- src/host/linalg.jl | 10 ++------ src/host/mapreduce.jl | 49 ++++++++++++++++++++++++++++-------- test/testsuite/reductions.jl | 48 +++++++++++++++++++++++++++++++++++ 5 files changed, 111 insertions(+), 30 deletions(-) diff --git a/Project.toml b/Project.toml index 1bd6c033..fd328d40 100644 --- a/Project.toml +++ b/Project.toml @@ -27,7 +27,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JLD2Ext = "JLD2" [compat] -AcceleratedKernels = "0.5" +AcceleratedKernels = "0.5.8" Adapt = "4.7.3" GPUArraysCore = "= 0.2.1" JLD2 = "0.4, 0.5, 0.6" diff --git a/src/host/indexing.jl b/src/host/indexing.jl index 6f5b1c6f..670cdd6b 100644 --- a/src/host/indexing.jl +++ b/src/host/indexing.jl @@ -212,9 +212,13 @@ _findlast_sentinel(::Type{I}) where {I <: Integer} = typemin(I) _findlast_sentinel(::Type{CartesianIndex{N}}) where {N} = CartesianIndex(ntuple(_ -> typemin(Int), Val(N))) _findfirst_op(::Type{<:Integer}) = min -_findfirst_op(::Type{<:CartesianIndex}) = (a, b) -> isless(b, a) ? b : a +_findfirst_op(::Type{<:CartesianIndex}) = _first_index _findlast_op(::Type{<:Integer}) = max -_findlast_op(::Type{<:CartesianIndex}) = (a, b) -> isless(a, b) ? b : a +_findlast_op(::Type{<:CartesianIndex}) = _last_index +_first_index(a, b) = isless(b, a) ? b : a +_last_index(a, b) = isless(a, b) ? b : a +_commutative(::Union{typeof(_first_index), typeof(_last_index)}, Ts::Type...) = + _commutes_on(CartesianIndex, Ts) struct _IndexIf{F, I} <: Function f::F @@ -280,16 +284,22 @@ struct _FindPair{F} end (p::_FindPair)(x, i) = (p.f(x), i) +# The pair with the greater value by `binop` (`isless` for `findmax`), or of equal values the one +# with the smaller index. It commutes where `binop` is a total order, as for real numbers. +struct _FindReduction{B} <: Function + binop::B +end +function (r::_FindReduction)(t1, t2) + (x, i), (y, j) = t1, t2 + r.binop(x, y) && return t2 + isequal(x, y) && return (x, min(i, j)) + return t1 +end +_commutative(::_FindReduction, Ts::Type...) = + _commutes_on(Tuple{_CommutativeReal, Integer}, Ts) -function findminmax(binop, f, A::AnyGPUArray; dims) - function reduction(t1, t2) - (x, i), (y, j) = t1, t2 - - binop(x, y) && return t2 - isequal(x, y) && return (x, min(i, j)) - return t1 - end +function findminmax(binop, f, A::AnyGPUArray; dims) if isempty(A) h = (binop === Base.isless ? findmax : findmin)(f, Array{eltype(A)}(undef, size(A)); dims) return dims === Colon() ? h : (copyto!(similar(A, eltype(h[1]), size(h[1])), h[1]), @@ -299,7 +309,7 @@ function findminmax(binop, f, A::AnyGPUArray; dims) # (along valid `dims`, a 0-dimensional array reduces nothing; Julia 1.10's `reduced_indices` # cannot take it) rdims = ndims(A) == 0 && !(dims isa Colon) && all(d -> d isa Integer && d >= 1, dims) ? () : dims - res = mapreduce(_FindPair(f), reduction, A, LinearIndices(A); dims=rdims) + res = mapreduce(_FindPair(f), _FindReduction(binop), A, LinearIndices(A); dims=rdims) I = keys(A) if dims === Colon() return (res[1], I[res[2]]) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index a198c3de..259eb1ed 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -242,10 +242,7 @@ function LinearAlgebra.istriu(A::AbstractGPUMatrix, k::Integer = 0) true end end - function reducer(a, b) - a && b - end - mapreduce(mapper, reducer, A, eachindex(IndexCartesian(), A); init=true) + mapreduce(mapper, &, A, eachindex(IndexCartesian(), A); init=true) end # check if lower triangular starting from the kth subdiagonal. @@ -258,10 +255,7 @@ function LinearAlgebra.istril(A::AbstractGPUMatrix, k::Integer = 0) true end end - function reducer(a, b) - a && b - end - mapreduce(mapper, reducer, A, eachindex(IndexCartesian(), A); init=true) + mapreduce(mapper, &, A, eachindex(IndexCartesian(), A); init=true) end # istriu/istril of triangular wrappers of the opposite triangularity: the generic diff --git a/src/host/mapreduce.jl b/src/host/mapreduce.jl index 2c29ad09..9c1b1810 100644 --- a/src/host/mapreduce.jl +++ b/src/host/mapreduce.jl @@ -11,8 +11,35 @@ const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcast.Broadcasted} # AcceleratedKernels' primitives, on the backend `backend` (a `Broadcasted` source may hold arrays # without one, such as `EachIndex`). The reference back-end (JLArrays) runs them on its storage # instead; these are internal, not an extension point for back-ends. -_ak_mapreduce(f, op, A; kwargs...) = AK.mapreduce(f, op, A; kwargs...) -_ak_mapreducedim!(f, op, R, A; kwargs...) = AK.mapreducedim!(f, op, R, A; kwargs...) +# +# Base's reductions combine the elements in their order, so `op` need not be commutative, while +# AcceleratedKernels' faster kernels combine them in any order: they are only used for operators +# known to commute on the types involved (see `_commutative`). +_ak_mapreduce(f, op, A; kwargs...) = + AK.mapreduce(f, op, A; commutative=_commutative(op, _mapped_eltype(f, A), _init_type(kwargs)), + kwargs...) +_ak_mapreducedim!(f, op, R, A; kwargs...) = + AK.mapreducedim!(f, op, R, A; commutative=_commutative(op, _mapped_eltype(f, A), eltype(R), + _init_type(kwargs)), + kwargs...) +_init_type(kwargs) = haskey(kwargs, :init) ? typeof(kwargs[:init]) : Union{} +_mapped_eltype(f, A) = Base.promote_op(f, _source_eltype(A)) + +# Whether `op` commutes on values of the types `Ts` (`Union{}` stands for no value): Base's +# arithmetic and bitwise operators on Base's number types, and this package's own commutative +# operators. Other types can define these operators without commuting (`*` of matrices), so +# anything else is reduced in order. +const _CommutativeReal = Union{Bool, Base.BitInteger, Base.IEEEFloat} +const _CommutativeNumber = Union{_CommutativeReal, Complex{<:_CommutativeReal}} +_commutes_on(::Type{S}, Ts) where {S} = all(T -> T <: S, Ts) +_commutative(op, Ts::Type...) = false +_commutative(::Union{typeof(+), typeof(*), typeof(Base.add_sum), typeof(Base.mul_prod)}, + Ts::Type...) = _commutes_on(_CommutativeNumber, Ts) +_commutative(::Union{typeof(max), typeof(min)}, Ts::Type...) = _commutes_on(_CommutativeReal, Ts) +_commutative(::Union{typeof(&), typeof(|), typeof(xor)}, Ts::Type...) = + _commutes_on(Union{Bool, Base.BitInteger}, Ts) +_commutative(::typeof(Base._extrema_rf), Ts::Type...) = + _commutes_on(Tuple{_CommutativeReal, _CommutativeReal}, Ts) # `Base.mapreducedim!` folds into `R`'s values Base.mapreducedim!(f, op, R::AnyGPUArray, A::AbstractArray) = @@ -259,13 +286,15 @@ function Base.:(==)(A::AnyGPUArray, B::AnyGPUArray) (; is_missing=false, is_equal=eq) end end - function reducer(a, b) - if a.is_missing || b.is_missing - (; is_missing=true, is_equal=#=don't care=#false) - else - (; is_missing=false, is_equal=a.is_equal & b.is_equal) - end - end - res = mapreduce(mapper, reducer, A, B; init=(; is_missing=false, is_equal=true)) + res = mapreduce(mapper, _equal_reducer, A, B; init=(; is_missing=false, is_equal=true)) res.is_missing ? missing : res.is_equal end +function _equal_reducer(a, b) + if a.is_missing || b.is_missing + (; is_missing=true, is_equal=#=don't care=#false) + else + (; is_missing=false, is_equal=a.is_equal & b.is_equal) + end +end +_commutative(::typeof(_equal_reducer), Ts::Type...) = + _commutes_on(@NamedTuple{is_missing::Bool, is_equal::Bool}, Ts) diff --git a/test/testsuite/reductions.jl b/test/testsuite/reductions.jl index dc9ebec0..f48690c4 100644 --- a/test/testsuite/reductions.jl +++ b/test/testsuite/reductions.jl @@ -268,6 +268,54 @@ end @test compare((A, B) -> isequal(A, B), AT, [missing], [missing]) end +# Associative operators that do not commute: the elements are combined in their order, as in Base +@testsuite "reductions/order" (AT, eltypes)->begin + # The first maximum and its index (#164 in AcceleratedKernels.jl): ties far apart, and in + # neighbouring blocks' or threads' elements + function first_max((xmax, imax), (x, i)) + xmax < x ? (x, i) : (xmax, imax) + end + for js in ((2, 163841), (2, 257), (300, 4097, 70_000)) + xs = zeros(Float32, 1_000_000) + xs[collect(js)] .= 1 + @test mapreduce(tuple, first_max, AT(xs), AT(collect(eachindex(xs))); + init=(-Inf32, -1)) == (1f0, first(js)) + end + + # The last non-zero value + last_nonzero(a, b) = iszero(b) ? a : b + @test compare(A -> reduce(last_nonzero, A), AT, rand(Int32(-2):Int32(2), 100_000)) + @test compare(A -> reduce(last_nonzero, A; dims=1, init=Int32(0)), AT, + rand(Int32(-2):Int32(2), 3000, 30)) + @test compare(A -> reduce(last_nonzero, A; dims=2, init=Int32(0)), AT, + rand(Int32(-2):Int32(2), 30, 3000)) + + # Products of 2x2 matrices over UInt32 (as tuples), random products of SL(2, Z) generators so + # that long products do not collapse to zero + matmul((a, b, c, d), (e, f, g, h)) = (a * e + b * g, a * f + b * h, c * e + d * g, c * f + d * h) + o, z, m = UInt32(1), UInt32(0), typemax(UInt32) + gens = ((o, o, z, o), (o, z, o, o), (o, m, z, o), (o, z, m, o)) + I2 = (o, z, z, o) + randmat() = foldl(matmul, rand(gens, 3); init=I2) + M = [randmat() for _ in 1:200_000] + @test compare(A -> reduce(matmul, A), AT, M) + init = randmat() + @test compare(A -> reduce(matmul, A; init), AT, M) + for (sz, dims) in (((40, 5000), 2), ((5000, 40), 1), ((300, 300), 1), ((300, 300), 2), + ((20, 30, 40), (1, 3))) + @test compare(A -> reduce(matmul, A; dims, init=I2), AT, + [randmat() for _ in CartesianIndices(sz)]) + end + + # Base's operators on Base's number types commute; on other types they need not + @test GPUArrays._commutative(+, Float32, Union{}) + @test GPUArrays._commutative(Base.add_sum, Int8, Int) + @test GPUArrays._commutative(max, Float64, Float64) + @test !GPUArrays._commutative(*, NTuple{4, UInt32}, Union{}) + @test !GPUArrays._commutative(+, Float32, Any) + @test !GPUArrays._commutative((a, b) -> a + b, Float32, Float32) +end + @testsuite "reductions/contract" (AT, eltypes)->begin M = rand(1:10, 10, 20)