add sigmoid and relu support

Author: chriselrod

Project.toml

@@ -6,7 +6,6 @@ version = "0.9.19"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
@@ -20,9 +19,10 @@ ArrayInterface = "2.14.12"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
 OffsetArrays = "1.4.1, 1.5"
-SLEEFPirates = "0.6.4"
+Requires = "1"
+SLEEFPirates = "0.6.6"
 UnPack = "1"
-VectorizationBase = "0.15.2"
+VectorizationBase = "0.15.3"
 julia = "1.5"
 
 [extras]

src/LoopVectorization.jl

@@ -37,7 +37,7 @@ using Requires
 
 export LowDimArray, stridedpointer,
     @avx, @_avx, *ˡ, _avx_!,
-    vmap, vmap!, vmapt, vmapt!, vmapnt, vmapnt!, vmapntt, vmapntt!, tanh_fast,
+    vmap, vmap!, vmapt, vmapt!, vmapnt, vmapnt!, vmapntt, vmapntt!, tanh_fast, sigmoid_fast,
     vfilter, vfilter!, vmapreduce, vreduce
 
 @inline unwrap(::Val{N}) where {N} = N
@@ -96,8 +96,12 @@ LoopVectorization
 include("precompile.jl")
 _precompile_()
 
+# import ChainRulesCore, ForwardDiff
+# include("vmap_grad.jl")
 function __init__()
-    @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" include("vmap_grad.jl")
+    @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" begin
+        @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include("vmap_grad.jl")
+    end
 end
 
 end # module

src/vmap_grad.jl

@@ -1,5 +1,4 @@
 
-using ForwardDiff
 using VectorizationBase: AbstractSIMD
 
 @generated function SLEEFPirates.tanh_fast(x::ForwardDiff.Dual{T,S,N}) where {T,S,N}
@@ -13,7 +12,46 @@ using VectorizationBase: AbstractSIMD
 end
 function ChainRulesCore.rrule(::typeof(tanh_fast), x)
     t = tanh_fast(x)
-    t, y -> (ChainRulesCore.Zero(), mul_fast(vfnmadd_fast(t, t, one(t)), y))
+    ∂ = let t = t
+        y -> (ChainRulesCore.Zero(), mul_fast(vfnmadd_fast(t, t, one(t)), y))
+    end
+    t, ∂
+end
+@generated function SLEEFPirates.sigmoid_fast(x::ForwardDiff.Dual{T,S,N}) where {T,S,N}
+    quote
+        $(Expr(:meta,:inline))
+        s = sigmoid_fast(x.value)
+        ∂s = vfnmadd_fast(s,s,s)
+        p = x.partials
+        ForwardDiff.Dual(s, ForwardDiff.Partials(Base.Cartesian.@ntuple $N n -> mul_fast(∂s, p[n])))
+    end
+end
+function ChainRulesCore.rrule(::typeof(sigmoid_fast), x)
+    s = sigmoid_fast(x)
+    ∂ = let s = s
+        y -> (ChainRulesCore.Zero(), mul_fast(vfnmadd_fast(s, s, s), y))
+    end
+    s, ∂
+end
+@generated function VectorizationBase.relu(x::ForwardDiff.Dual{T,S,N}) where {T,S,N}
+    quote
+        $(Expr(:meta,:inline))
+        v = x.value
+        z = zero(v)
+        cmp = v < z
+        r = ifelse(cmp, z, v)
+        p = x.partials
+        ForwardDiff.Dual(r, ForwardDiff.Partials(Base.Cartesian.@ntuple $N n -> ifelse(cmp, z, p[n])))
+    end
+end
+function ChainRulesCore.rrule(::typeof(relu), v)
+    z = zero(v)
+    cmp = v < z
+    r = ifelse(cmp, z, v)
+    ∂ = let cmp = cmp
+        y -> (ChainRulesCore.Zero(), ifelse(cmp, zero(y), y))
+    end
+    r, ∂
 end
 
 
@@ -33,16 +71,6 @@ end
     q
 end
 
-@generated function dual_store!(∂p::Tuple{Vararg{AbstractStridedPointer,A}}, p::AbstractStridedPointer, ∂v, im::Vararg{Any,N}) where {A,N}
-    quote
-        $(Expr(:meta,:inline))
-        v = ∂v.value
-        ∂ = ∂v.partials
-        VectorizationBase.vnoaliasstore!(p, v, im...)
-        Base.Cartesian.@nexprs $A a -> VectorizationBase.vnoaliasstore!(∂p[a], ∂[a], im...)
-        nothing
-    end
-end
 @generated function dual_store!(∂p::Tuple{Vararg{AbstractStridedPointer,A}}, p::AbstractStridedPointer, ∂v, im::Vararg{Any,N}) where {A,N}
     quote
         $(Expr(:meta,:inline))

test/zygote.jl

@@ -1,22 +1,29 @@
 using Zygote
-zygotefun(a,b,c) = log(tanh_fast(a) + 1 + exp(b - c))
+zygotefun2(a,b) = sigmoid_fast(a + b) + LoopVectorization.relu(a * b) * tanh(a - b)
+zygotefun3(a,b,c) = log(tanh_fast(a*b) + 1 + exp(b - c))
 @testset "Zygote" begin
     @show @__LINE__
 
     for T ∈ (Float32,Float64)
         x = randn(T, 217); y = randn(T, 217); z = randn(T, 217);
-
+        w = randn(T, 217)
         # Test 1-arg and `tanh_fast` vs `tanh`
-        gtref = gradient(x -> sum(map(tanh, x)), x);
-        gtlv = gradient(x -> sum(vmap(tanh_fast, x)), x);
+        gref1 = gradient(x -> dot(w,  map(tanh,      x)), x);
+        glv_1 = gradient(x -> dot(w, vmap(tanh_fast, x)), x);
+
+        @test  only(gref1) ≈ only(glv_1)
+
+        # Test 2 arguments
+        # Test 3 arguments
+        gref2 = gradient(xyz -> dot(w,  map(zygotefun2, xyz...)), (x, y));
+        g_lv2 = gradient(xyz -> dot(w, vmap(zygotefun2, xyz...)), (x, y));
 
-        @test  only(gtref) ≈ only(gtlv)
+        @test all(map(≈, only(gref2), only(g_lv2)))
 
-        # Test multiple arguments
-        gref = gradient(xyz -> sum(map(zygotefun, xyz...)), (x, y, z));
-        glv = gradient(xyz -> sum(vmap(zygotefun, xyz...)), (x, y, z));
+        gref3 = gradient(xyz -> dot(w,  map(zygotefun3, xyz...)), (x, y, z));
+        g_lv3 = gradient(xyz -> dot(w, vmap(zygotefun3, xyz...)), (x, y, z));
 
-        @test all(map(≈, only(gref), only(glv)))
+        @test all(map(≈, only(gref3), only(g_lv3)))
     end
 end