Simplify fastpow implementation via T output

Author: ChrisRackauckas

ext/DiffEqBaseEnzymeExt.jl

@@ -53,46 +53,6 @@ function Enzyme.EnzymeRules.reverse(config::Enzyme.EnzymeRules.RevConfigWidth{1}
     return ntuple(_ -> nothing, Val(length(args) + 4))
 end
 
-function Enzyme.EnzymeRules.forward(func::Const{typeof(DiffEqBase.fastpow)},
-        RT::Type{<:Union{Duplicated, DuplicatedNoNeed}},
-        _x::Union{Const, Duplicated}, _y::Union{Const, Duplicated})
-    x = _x.val
-    y = _y.val
-    ret = func.val(x, y)
-    if !(_x isa Const) 
-        dxval = _x.dval * y * (fastpow(x,y - 1))  
-    else 
-        dxval = make_zero(_x.val)
-    end
-    if !(_y isa Const)  
-        dyval = x isa Real && x<=0 ? Base.oftype(float(x), NaN) : _y.dval*(fastpow(x,y))*log(x)
-    else 
-        dyval = make_zero(_y.val)
-    end 
-    if RT <: DuplicatedNoNeed
-        return Float32(dxval + dyval)
-    else
-        return Duplicated(ret, Float32(dxval + dyval))
-    end
-end
-
-function EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.ConfigWidth{1}, 
-        func::Const{typeof(fastpow)}, ::Type{<:Active}, x::Active, y::Active)
-    if EnzymeRules.needs_primal(config)
-        primal = func.val(x.val, y.val)
-    else
-        primal = nothing
-    end
-    return EnzymeRules.AugmentedReturn(primal, nothing, nothing)
-end
-
-function EnzymeRules.reverse(config::Enzyme.EnzymeRules.ConfigWidth{1}, 
-        func::Const{typeof(DiffEqBase.fastpow)}, dret::Active, tape, _x::Active, _y::Active)
-    x = _x.val
-    y = _y.val
-    dxval =  dret.val * y * (fastpow(x,y - 1))
-    dyval = x isa Real && x<=0 ? Base.oftype(float(x), NaN) : dret.val * (fastpow(x,y))*log(x)
-    return (dxval, dyval)
-end
+Enzyme.Compiler.known_ops[typeof(DiffEqBase.fastpow)] = (:pow, 2, nothing)
 
 end 

src/fastpow.jl

@@ -51,60 +51,19 @@ const EXP2FT = (Float32(0x1.6a09e667f3bcdp-1),
     Float32(0x1.3dea64c123422p+0),
     Float32(0x1.4bfdad5362a27p+0),
     Float32(0x1.5ab07dd485429p+0))
-@inline function _exp2(x::Float32)
-    TBLBITS = UInt32(4)
-    TBLSIZE = UInt32(1 << TBLBITS)
 
-    redux = Float32(0x1.8p23) / TBLSIZE
-    P1 = Float32(0x1.62e430p-1)
-    P2 = Float32(0x1.ebfbe0p-3)
-    P3 = Float32(0x1.c6b348p-5)
-    P4 = Float32(0x1.3b2c9cp-7)
-
-    # Reduce x, computing z, i0, and k.
-    t::Float32 = x + redux
-    i0 = reinterpret(UInt32, t)
-    i0 += TBLSIZE ÷ UInt32(2)
-    k::UInt32 = unsafe_trunc(UInt32, (i0 >> TBLBITS) << 20)
-    i0 &= TBLSIZE - UInt32(1)
-    t -= redux
-    z = x - t
-    twopk = Float32(reinterpret(Float64, UInt64(0x3ff00000 + k) << 32))
-
-    # Compute r = exp2(y) = exp2ft[i0] * p(z).
-    tv = EXP2FT[i0 + UInt32(1)]
-    u = tv * z
-    tv = tv + u * (P1 + z * P2) + u * (z * z) * (P3 + z * P4)
-
-    # Scale by 2**(k>>20)
-    return tv * twopk
-end
-
-if VERSION < v"1.7.0"
-    """
-        fastpow(x::Real, y::Real) -> Float32
-    """
-    @inline function fastpow(x::Real, y::Real)
-        if iszero(x)
-            return 0.0f0
-        elseif isinf(x) && isinf(y)
-            return Float32(Inf)
-        else
-            return _exp2(convert(Float32, y) * fastlog2(convert(Float32, x)))
-        end
-    end
-else
-    """
-        fastpow(x::Real, y::Real) -> Float32
-    """
-    @inline function fastpow(x::Real, y::Real)
-        if iszero(x)
-            return 0.0f0
-        elseif isinf(x) && isinf(y)
-            return Float32(Inf)
-        else
-            return @fastmath exp2(convert(Float32, y) * fastlog2(convert(Float32, x)))
-        end
+"""
+    fastpow(x::T, y::T) where {T} -> T
+    Trips through Float32 for performance.
+"""
+@inline function fastpow(x::T, y::T) where {T}
+    if iszero(x)
+        return zero(T)
+    elseif isinf(x) && isinf(y)
+        return convert(T,Inf)
+    else
+        return convert(T,@fastmath exp2(convert(Float32, y) * fastlog2(convert(Float32, x))))
     end
 end
+
 @inline fastpow(x, y) = x^y