Added vmap and vmap!

Author: chriselrod

src/LoopVectorization.jl

@@ -9,7 +9,8 @@ using MacroTools: prewalk, postwalk
 
 
 export LowDimArray, stridedpointer, vectorizable,
-    @vectorize, @vvectorize, @avx, ∗
+    @vectorize, @vvectorize, @avx, ∗,
+    vmap, vmap!
 
 function isdense end #
 
@@ -904,6 +905,7 @@ include("broadcast.jl")
 include("determinestrategy.jl")
 include("lowering.jl")
 include("constructors.jl")
+include("map.jl")
 include("precompile.jl")
 _precompile_()
 

src/map.jl

@@ -0,0 +1,31 @@
+
+
+
+@generated function vmap!(f::F, dest::AbstractArray{T}, args::Vararg{<:AbstractArray,N}) where {F,T,N}
+    W, Wshift = VectorizationBase.pick_vector_width_shift(T)
+    val = Expr(:curly, :Val, W)
+    q = Expr(:block, Expr(:(=), :M, Expr(:call, :length, :dest)), Expr(:(=), :vdest, Expr(:call, :vectorizable, :dest)), Expr(:(=), :m, 0))
+    fcall = Expr(:call, :f)
+    loopbody = Expr(:block, Expr(:call, :vstore!, :vdest, fcall, :m), Expr(:(+=), :m, W))
+    fcallmask = Expr(:call, :f)
+    bodymask = Expr(:block, Expr(:(=), :__mask__, Expr(:call, :mask, val, Expr(:call, :&, :M, W-1))), Expr(:call, :vstore!, :vdest, fcallmask, :m, :__mask__))
+    for n ∈ 1:N
+        arg_n = Symbol(:varg_,n)
+        push!(q.args, Expr(:(=), arg_n, Expr(:call, :vectorizable, Expr(:ref, :args, n))))
+        push!(fcall.args, Expr(:call, :vload, Expr(:call, val), arg_n, :m))
+        push!(fcallmask.args, Expr(:call, :vload, Expr(:call, val), arg_n, :m, :__mask__))
+    end
+    loop = Expr(:for, Expr(:(=), :_, Expr(:call, :(:), 0, Expr(:call, :-, Expr(:call, :(>>>), :M, Wshift), 1))), loopbody)
+    push!(q.args, loop)
+    ifmask = Expr(:if, Expr(:call, :(!=), :m, :M), bodymask)
+    push!(q.args, ifmask)
+    push!(q.args, :dest)
+    q
+end
+function vmap(f::F, args...) where {F}
+    T = Base._return_type(f, Base.Broadcast.eltypes(args))
+    dest = similar(first(args), T)
+    vmap!(f, dest, args...)
+end
+
+

test/runtests.jl

@@ -358,4 +358,16 @@ end
         @test D1 ≈ D2
     end
 end
+
+@testset "map" begin
+    foo(x, y) = exp(x) - sin(y)
+    N = 37
+    for T ∈ (Float32,Float64)
+        a = rand(T, N); b = rand(T, N)
+        c1 = map(foo, a, b)
+        c2 = vmap(foo, a, b)
+        @test c1 ≈ c2
+    end
+end
+
 end