Reduce with operators not known to commute in element order

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"

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]])

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

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)

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)