Merge pull request #17 from MasonProtter/_avx

Refactor @avx with generated functions

Author: chriselrod

Manifest.toml

@@ -61,9 +61,9 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [[SIMDPirates]]
 deps = ["MacroTools", "VectorizationBase"]
-git-tree-sha1 = "910193d289b41e570118c4e444f0c05cc700a2f7"
+git-tree-sha1 = "500294a8b1001bdda2483fc6d675956798ad8764"
 uuid = "21efa798-c60a-11e8-04d3-e1a92915a26a"
-version = "0.1.5"
+version = "0.1.6"
 
 [[SLEEFPirates]]
 deps = ["SIMDPirates", "VectorizationBase"]

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.3.6"
+version = "0.3.7"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -14,7 +14,7 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 [compat]
 MacroTools = "0.5"
 Parameters = "0.12.0"
-SIMDPirates = "0.1.5"
+SIMDPirates = "0.1.6"
 SLEEFPirates = "0.1.3"
 VectorizationBase = "0.1.9"
 julia = "1.3.0"

src/LoopVectorization.jl

@@ -22,6 +22,10 @@ include("determinestrategy.jl")
 include("lowering.jl")
 include("constructors.jl")
 include("map.jl")
+include("_avx.jl")
+
+export @_avx, _avx
+
 # include("precompile.jl")
 # _precompile_()
 

src/_avx.jl

@@ -0,0 +1,78 @@
+using MacroTools: postwalk, prewalk
+using LoopVectorization: LoopVectorization, LoopSet, lower
+
+#----------------------------------------------------------------------------------------------------
+
+struct Ex{T, Tup} end
+
+function to_type(ex::Expr)
+    Ex{ex.head, Tuple{to_type.(ex.args)...}}
+end
+
+to_type(x) = x
+to_type(::LineNumberNode) = nothing
+
+#----------------------------------------------------------------------------------------------------
+
+to_expr(ex::Type{Ex{Head, Tup}}) where {Head, Tup} = Expr(Head, (to_expr(x) for x in Tup.parameters)...)
+to_expr(x) = x
+
+#----------------------------------------------------------------------------------------------------
+
+function find_vars_and_gensym!(ex::Expr, vars::Set{Symbol}, ivars::Vector{Symbol})
+    if ex.head == :(=) && ex.args[1] isa Symbol
+        push!(ivars, ex.args[1])
+    elseif ex.head == :call
+        push!(ivars, ex.args[1])
+    end
+    ex
+end
+
+function find_vars_and_gensym!(x::Symbol, vars::Set{Symbol}, ivars::Vector{Symbol})
+    if (x ∉ vars) && (x ∉ ivars)
+        push!(vars, x)
+        x
+    else
+        x
+    end    
+end
+
+find_vars_and_gensym!(x, vars::Set{Symbol}, ivars::Vector{Symbol}) = x
+
+#----------------------------------------------------------------------------------------------------
+
+nt(keys, vals) = NamedTuple{keys, typeof(vals)}(vals)
+
+macro _avx(ex)
+    D     = Set{Symbol}()
+    ivars = Symbol[]
+
+    gex = prewalk(x -> find_vars_and_gensym!(x, D, ivars), ex)
+
+    type_ex = to_type(gex)
+
+    tvars  = Tuple(D)
+
+    quote
+        kwargs = LoopVectorization.nt($(QuoteNode(tvars)), $(Expr(:tuple, tvars...)))
+        $(Expr(:tuple, tvars...)) = LoopVectorization._avx($(QuoteNode(type_ex)), kwargs)
+        # LoopVectorization._avx($(QuoteNode(type_ex)), kwargs) # comment out the above line, uncomment this one, and get rid of the `@generated` on _avx to see the function body.
+    end |> esc
+end
+
+@generated function _avx(::Type{ex_t}, var_nt::NamedTuple{keys, var_types}) where {ex_t <: Ex, keys, var_types}
+    ex = to_expr(ex_t)
+    
+    var_defs = Expr(:block, )
+    for k in keys 
+        push!(var_defs.args, :($k = var_nt[$(QuoteNode(k))]))
+    end
+
+    quote
+        $(Expr(:meta,:inline))
+        $var_defs
+        $(lower(LoopSet(ex)))
+        $(Expr(:tuple, keys...))
+        #$(Expr(:tuple, (:($(keys[i]) :: $(var_types.parameters[i])) for i in eachindex(keys))...))
+    end
+end