ForwardDiffExt: switched to NNlib activation functions

Author: mxpoch

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
-authors = ["Chris Elrod <elrodc@gmail.com>"]
 version = "0.12.172"
+authors = ["Chris Elrod <elrodc@gmail.com>"]
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -30,9 +30,10 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
 [extensions]
-ForwardDiffExt = ["ChainRulesCore", "ForwardDiff"]
+ForwardDiffExt = ["ChainRulesCore", "ForwardDiff", "NNlib"]
 SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -46,6 +47,7 @@ HostCPUFeatures = "0.1.10"
 IfElse = "0.1"
 LayoutPointers = "0.1.11"
 LinearAlgebra = "1"
+NNlib = "0.9.31"
 OffsetArrays = "1.4.1"
 PolyesterWeave = "0.1.10, 0.2"
 PrecompileTools = "1"
@@ -57,4 +59,8 @@ StaticArrayInterface = "1"
 ThreadingUtilities = "0.5"
 UnPack = "1"
 VectorizationBase = "0.21.72"
-julia = "1.6"
+julia = "1.10"
+
+[extras]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"

ext/ForwardDiffExt.jl

@@ -1,14 +1,14 @@
 module ForwardDiffExt
 import ForwardDiff, ChainRulesCore
-using LoopVectorization, VectorizationBase, SLEEFPirates, ForwardDiff
+using LoopVectorization, VectorizationBase, SLEEFPirates, ForwardDiff, NNlib
+using SLEEFPirates: tanh_fast, sigmoid_fast
 
 import IfElse: ifelse
 using VectorizationBase: AbstractSIMD, AbstractMask, zero_offsets
 
 using LoopVectorization:
   AbstractSIMD,
   AbstractStridedPointer,
-  relu,
   vmap,
   VectorizationBase,
   vmapt,
@@ -140,7 +140,8 @@ end
     )
   end
 end
-@generated function VectorizationBase.relu(
+
+@generated function NNlib.relu(
   x::ForwardDiff.Dual{T,S,N}
 ) where {T,S,N}
   quote
@@ -157,6 +158,27 @@ end
   end
 end
 
+@generated function NNlib.leakyrelu(
+  x::ForwardDiff.Dual{T,S,N},
+  a = 0.01
+) where {T,S,N}
+  quote
+    $(Expr(:meta, :inline))
+    v = x.value
+    z = zero(v)
+
+    α = convert(typeof(v), a)
+    cmp = v < z
+    r = ifelse(cmp, α * v, v)
+    p = x.partials
+    ForwardDiff.Dual{T}(
+      r,
+      ForwardDiff.Partials(Base.Cartesian.@ntuple $N n -> ifelse(cmp, α * p[n], p[n]))
+    )
+  end
+end
+
+
 @generated function _ifelse(
   m::Union{AbstractMask,VecUnroll{<:Any,<:Any,Bit,<:AbstractMask}},
   x::ForwardDiff.Dual{TAG,V,P},

test/forwarddiffext.jl

@@ -16,21 +16,6 @@ function tovec(x::ForwardDiff.Dual{T,V,N}) where {T,V,N}
   return ret
 end
 
-if LoopVectorization.ifelse !== Base.ifelse
-  @inline function NNlib.leakyrelu(
-    x::LoopVectorization.AbstractSIMD,
-    a = NNlib.oftf(x, NNlib.leakyrelu_a),
-  )
-    LoopVectorization.ifelse(x > zero(x), float(x), NNlib.oftf(x, a * x))  # max(a*x, x) is 3x slower
-  end
-  @inline function NNlib.leakyrelu(
-    x::ForwardDiff.Dual{<:Any,<:LoopVectorization.AbstractSIMD},
-    a = NNlib.oftf(x, NNlib.leakyrelu_a),
-  )
-    LoopVectorization.ifelse(x > zero(x), float(x), NNlib.oftf(x, a * x))  # max(a*x, x) is 3x slower
-  end
-end
-
 vx0 = randnvec()
 vx1 = randnvec()
 vx2 = randnvec()
@@ -50,3 +35,8 @@ vud = ForwardDiff.Dual(vu0, vu1, vu2)
       reinterpret(Float64, NNlib.leakyrelu.(tovec(vd0)))
 @test reinterpret(Float64, tovec(NNlib.leakyrelu(vud))) ≈
       reinterpret(Float64, NNlib.leakyrelu.(tovec(vud)))
+
+@test reinterpret(Float64, tovec(NNlib.relu(vd0))) ≈
+      reinterpret(Float64, NNlib.relu.(tovec(vd0)))
+@test reinterpret(Float64, tovec(NNlib.relu(vud))) ≈
+      reinterpret(Float64, NNlib.relu.(tovec(vud)))