Define neutral elements for the @fastmath reduction operators

Project.toml

@@ -29,7 +29,7 @@ JLD2Ext = "JLD2"
 [compat]
 AcceleratedKernels = "0.5.8"
 Adapt = "4.7.3"
-GPUArraysCore = "= 0.2.1"
+GPUArraysCore = "= 0.2.2"
 JLD2 = "0.4, 0.5, 0.6"
 KernelAbstractions = "0.9.43, 0.10"
 LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10"

lib/GPUArraysCore/Project.toml

@@ -1,7 +1,7 @@
 name = "GPUArraysCore"
 uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/GPUArraysCore/src/GPUArraysCore.jl

@@ -221,8 +221,9 @@ seed partial results, e.g. `neutral_element(+, Float32) === 0.0f0` and
 
 Methods exist for the operators that Base's reductions use: `+`, `*`, `&`, `|`, `⊻`, `min`,
 `max`, `Base.add_sum`, `Base.mul_prod` and `Base._extrema_rf`, and `Base.and_all` and
-`Base.or_any` on the Julia versions that define them. For any other operator this throws an
-error; packages can add a method for an operator whose neutral element they know.
+`Base.or_any` on the Julia versions that define them. The same goes for the `Base.FastMath`
+versions of `+`, `*`, `min` and `max`, which `@fastmath` substitutes. For any other operator
+this throws an error; packages can add a method for an operator whose neutral element they know.
 """
 neutral_element(op, T) =
     error("""The neutral element of the operator `$op` for elements of type `$T` is unknown.
@@ -237,6 +238,12 @@ neutral_element(::typeof(Base.:(*)), T) = one(T)
 neutral_element(::typeof(Base.mul_prod), T) = one(T)
 neutral_element(::typeof(Base.min), T) = typemax(T)
 neutral_element(::typeof(Base.max), T) = typemin(T)
+# the operators that `@fastmath` substitutes: without a neutral element, reductions with them take
+# a slower generic path than reductions with the operators they replace
+neutral_element(::typeof(Base.FastMath.add_fast), T) = zero(T)
+neutral_element(::typeof(Base.FastMath.mul_fast), T) = one(T)
+neutral_element(::typeof(Base.FastMath.min_fast), T) = typemax(T)
+neutral_element(::typeof(Base.FastMath.max_fast), T) = typemin(T)
 neutral_element(::typeof(Base._extrema_rf), ::Type{<:NTuple{2,T}}) where {T} = typemax(T), typemin(T)
 @static if isdefined(Base, :and_all) # VERSION >~ v"1.13-"
     neutral_element(::typeof(Base.:(and_all)), T) = ~zero(T)

src/host/mapreduce.jl

@@ -26,16 +26,20 @@ _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
+# arithmetic and bitwise operators on Base's number types (including the `@fastmath` versions,
+# which may ignore NaNs and the sign of zero anyway), 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)},
+_commutative(::Union{typeof(+), typeof(*), typeof(Base.add_sum), typeof(Base.mul_prod),
+                     typeof(Base.FastMath.add_fast), typeof(Base.FastMath.mul_fast)},
              Ts::Type...) = _commutes_on(_CommutativeNumber, Ts)
-_commutative(::Union{typeof(max), typeof(min)}, Ts::Type...) = _commutes_on(_CommutativeReal, Ts)
+_commutative(::Union{typeof(max), typeof(min),
+                     typeof(Base.FastMath.max_fast), typeof(Base.FastMath.min_fast)},
+             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...) =

test/testsuite/reductions.jl

@@ -311,6 +311,9 @@ end
     @test GPUArrays._commutative(+, Float32, Union{})
     @test GPUArrays._commutative(Base.add_sum, Int8, Int)
     @test GPUArrays._commutative(max, Float64, Float64)
+    @test GPUArrays._commutative(Base.FastMath.add_fast, Float32, Float32)
+    @test GPUArrays._commutative(Base.FastMath.max_fast, Float16, Union{})
+    @test !GPUArrays._commutative(Base.FastMath.mul_fast, NTuple{4, UInt32}, Union{})
     @test !GPUArrays._commutative(*, NTuple{4, UInt32}, Union{})
     @test !GPUArrays._commutative(+, Float32, Any)
     @test !GPUArrays._commutative((a, b) -> a + b, Float32, Float32)
@@ -512,9 +515,23 @@ end
     @test GPUArrays.neutral_element(max, Float64) === -Inf
     @test GPUArrays.neutral_element(&, UInt8) === 0xff
     @test GPUArrays.neutral_element(Base._extrema_rf, NTuple{2,Int8}) === (typemax(Int8), typemin(Int8))
+    # ... including the operators that `@fastmath` substitutes
+    @test GPUArrays.neutral_element(Base.FastMath.add_fast, Float32) === 0.0f0
+    @test GPUArrays.neutral_element(Base.FastMath.mul_fast, Int32) === Int32(1)
+    @test GPUArrays.neutral_element(Base.FastMath.min_fast, Int16) === typemax(Int16)
+    @test GPUArrays.neutral_element(Base.FastMath.max_fast, Float64) === -Inf
     @test_throws ErrorException GPUArrays.neutral_element((x, y) -> x, Int)
     if isdefined(Base, :and_all)    # Julia 1.13
         @test GPUArrays.neutral_element(Base.and_all, Bool) === true
         @test GPUArrays.neutral_element(Base.or_any, Bool) === false
     end
+
+    @testset "@fastmath $ET" for ET in eltypes
+        ET <: Real || continue
+        x = rand(ET(1):ET(2), 10, 10)
+        @test compare(A->@fastmath(reduce(max, A; dims=1, init=typemin(ET))), AT, x)
+        @test compare(A->@fastmath(reduce(min, A; init=typemax(ET))), AT, x)
+        @test compare(A->@fastmath(reduce(+, A; dims=2, init=zero(ET))), AT, x)
+        @test compare(A->@fastmath(reduce(*, A; dims=1, init=one(ET))), AT, x)
+    end
 end