Add Zygote support for vmap.

Author: chriselrod

Project.toml

@@ -1,14 +1,16 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.9.18"
+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"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SLEEFPirates = "476501e8-09a2-5ece-8869-fb82de89a1fa"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"

src/LoopVectorization.jl

@@ -32,11 +32,12 @@ using ArrayInterface
 using ArrayInterface: OptionallyStaticUnitRange, Zero, One#, static_length
 const Static = ArrayInterface.StaticInt
 
+using Requires
 
 
 export LowDimArray, stridedpointer,
     @avx, @_avx, *ˡ, _avx_!,
-    vmap, vmap!, vmapt, vmapt!, vmapnt, vmapnt!, vmapntt, vmapntt!,
+    vmap, vmap!, vmapt, vmapt!, vmapnt, vmapnt!, vmapntt, vmapntt!, tanh_fast,
     vfilter, vfilter!, vmapreduce, vreduce
 
 @inline unwrap(::Val{N}) where {N} = N
@@ -45,7 +46,6 @@ export LowDimArray, stridedpointer,
 
 const VECTORWIDTHSYMBOL, ELTYPESYMBOL = Symbol("##Wvecwidth##"), Symbol("##Tloopeltype##")
 
-
 include("vectorizationbase_compat/contract_pass.jl")
 include("vectorizationbase_compat/subsetview.jl")
 include("closeopen.jl")
@@ -96,4 +96,8 @@ LoopVectorization
 include("precompile.jl")
 _precompile_()
 
+function __init__()
+    @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" include("vmap_grad.jl")
+end
+
 end # module

src/map.jl

@@ -47,33 +47,40 @@ function vmap_singlethread!(
     V = VectorizationBase.pick_vector_width_val(T)
     W = unwrap(V)
     st = VectorizationBase.static_sizeof(T)
-    zero_index = MM{W}(Static(0), st)
-    while i < N - ((W << 2) - 1)
+    UNROLL = 4
+    LOG2UNROLL = 2
+    while i < N - ((W << LOG2UNROLL) - 1)
 
-        # vstore!(stridedpointer(B), VectorizationBase.VecUnroll((v1,v2,v3)), VectorizationBase.Unroll{AU,1,3,AV,W64,zero(UInt)}((i, j, k)))
-        # vload(stridedpointer(B), VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((i,)))
-        
-        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((i,))
+        index = VectorizationBase.Unroll{1,1,UNROLL,1,W,0x0000000000000000}((i,))
         v = f(vload.(ptrargs, index)...)
         if NonTemporal
             vstorent!(ptry, v, index)
         else
             vnoaliasstore!(ptry, v, index)
         end
-        i = vadd_fast(i, 4W)
+        i = vadd_fast(i, StaticInt{UNROLL}() * W)
     end
-    while i < N - (W - 1) # stops at 16 when
-        vᵣ = f(vload.(ptrargs, ((MM{W}(i),),))...)
-        if NonTemporal
-            vstorent!(ptry, vᵣ, (MM{W}(i),))
-        else
-            vnoaliasstore!(ptry, vᵣ, (MM{W}(i),))
+    if Base.libllvm_version ≥ v"11"
+        Nm1 = vsub_fast(N, 1)
+        while i < N # stops at 16 when
+            m = mask(V, i, Nm1)
+            vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(i),),), m)...), (MM{W}(i,),), m)
+            i = vadd_fast(i, W)
+        end
+    else
+        while i < N - (W - 1) # stops at 16 when
+            vᵣ = f(vload.(ptrargs, ((MM{W}(i),),))...)
+            if NonTemporal
+                vstorent!(ptry, vᵣ, (MM{W}(i),))
+            else
+                vnoaliasstore!(ptry, vᵣ, (MM{W}(i),))
+            end
+            i = vadd_fast(i, W)
+        end
+        if i < N
+            m = mask(T, N & (W - 1))
+            vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(i),),), m)...), (MM{W}(i,),), m)
         end
-        i = vadd_fast(i, W)
-    end
-    if i < N
-        m = mask(T, N & (W - 1))
-        vnoaliasstore!(ptry, f(vload.(ptrargs, ((MM{W}(i),),), m)...), (MM{W}(i,),), m)
     end
     y
 end

src/vmap_grad.jl

@@ -0,0 +1,119 @@
+
+using ForwardDiff
+using VectorizationBase: AbstractSIMD
+
+@generated function SLEEFPirates.tanh_fast(x::ForwardDiff.Dual{T,S,N}) where {T,S,N}
+    quote
+        $(Expr(:meta,:inline))
+        t = tanh_fast(x.value)
+        ∂t = vfnmadd_fast(t, t, one(S))
+        p = x.partials
+        ForwardDiff.Dual(t, ForwardDiff.Partials(Base.Cartesian.@ntuple $N n -> mul_fast(∂t, p[n])))
+    end
+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))
+end
+
+
+@generated function init_dual(v::Tuple{Vararg{AbstractSIMD,A}}) where {A}
+    res = Expr(:tuple)
+    q = Expr(:block, Expr(:meta,:inline))
+    for a ∈ 1:A
+        v_a = Symbol(:v_,a)
+        push!(q.args, Expr(:(=), v_a, Expr(:ref, :v, a)))
+        partials = Expr(:tuple)
+        for i ∈ 1:A
+            push!(partials.args, Expr(:call, i == a ? :one : :zero, v_a))
+        end
+        push!(res.args, :(ForwardDiff.Dual($v_a, ForwardDiff.Partials($partials))))
+    end
+    push!(q.args, res)
+    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))
+        v = ∂v.value
+        ∂ = ∂v.partials
+        VectorizationBase.vnoaliasstore!(p, v, im...)
+        Base.Cartesian.@nexprs $A a -> VectorizationBase.vnoaliasstore!(∂p[a], ∂[a], im...)
+        nothing
+    end
+end
+
+function ∂vmap_singlethread!(
+    f::F, ∂y::Tuple{Vararg{DenseArray{T},A}}, y::DenseArray{T},
+    args::Vararg{<:DenseArray{<:Base.HWReal},A}
+) where {F,T <: Base.HWReal, A}
+    N = length(y)
+    ptry = VectorizationBase.zero_offsets(stridedpointer(y))
+    ptrargs = VectorizationBase.zero_offsets.(stridedpointer.(args))
+    ptr∂y = VectorizationBase.zero_offsets.(stridedpointer.(∂y))
+    
+    i = 0
+    V = VectorizationBase.pick_vector_width_val(T)
+    W = Int(V)
+    st = VectorizationBase.static_sizeof(T)
+    zero_index = MM{W}(StaticInt(0), st)
+    while i < N - ((W << 2) - 1)
+        index = VectorizationBase.Unroll{1,1,4,1,W,0x0000000000000000}((i,))
+        v = f(init_dual(vload.(ptrargs, index))...)
+        dual_store!(ptr∂y, ptry, v, index)
+        i = vadd_fast(i, 4W)
+    end
+    while i < N - (W - 1) 
+        vᵣ = f(init_dual(vload.(ptrargs, ((MM{W}(i),),)))...)
+        dual_store!(ptr∂y, ptry, vᵣ, (MM{W}(i),))
+        i = vadd_fast(i, W)
+    end
+    if i < N
+        m = mask(T, N & (W - 1))
+        dual_store!(ptr∂y, ptry, f(init_dual(vload.(ptrargs, ((MM{W}(i),),), m))...), (MM{W}(i,),), m)
+    end
+    nothing
+end
+
+
+struct SIMDMapBack{K,T<:Tuple{Vararg{Any,K}}}
+    jacs::T
+end
+@generated function (b::SIMDMapBack{K,T})(Δ::A) where {K,T,A}
+    preloop = Expr(:block, :(jacs = b.jacs))
+    loop_body = Expr(:block, :(Δᵢ = Δ[i]))
+    ret = Expr(:tuple, ChainRulesCore.Zero(), ChainRulesCore.Zero())
+    for k ∈1:K
+        jₖ = Symbol(:j_, k)
+        push!(preloop.args, :($jₖ = jacs[$k]))
+        push!(loop_body.args, :($jₖ[i] *= Δᵢ))
+        push!(ret.args, jₖ)
+    end
+    quote
+        $preloop
+        @avx for i ∈ eachindex(Δ)
+            $loop_body
+        end
+        $ret
+    end
+end
+
+function ChainRulesCore.rrule(::typeof(vmap), f::F, args::Vararg{Any,K}) where {F,K}
+    out = similar(first(args))
+    jacs = map(similar, args)
+    ∂vmap_singlethread!(f, jacs, out, args...)
+    out, SIMDMapBack(jacs)
+end
+
+

test/Project.toml

@@ -6,3 +6,4 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

test/runtests.jl

@@ -27,6 +27,8 @@ const START_TIME = time()
 
     @time include("check_empty.jl")
 
+    @time include("zygote.jl")
+
     @time include("offsetarrays.jl")
 
     @time include("tensors.jl")

test/zygote.jl

@@ -0,0 +1,23 @@
+using Zygote
+zygotefun(a,b,c) = log(tanh_fast(a) + 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);
+
+        # 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);
+
+        @test  only(gtref) ≈ only(gtlv)
+
+        # Test multiple arguments
+        gref = gradient(xyz -> sum(map(zygotefun, xyz...)), (x, y, z));
+        glv = gradient(xyz -> sum(vmap(zygotefun, xyz...)), (x, y, z));
+
+        @test all(map(≈, only(gref), only(glv)))
+    end
+end
+
+