@nospecialize a few type unstable functions and reduce to a single method to avoid dynamic dispatches.

Author: chriselrod

src/_avx.jl

@@ -5,12 +5,18 @@ using LoopVectorization: LoopVectorization, LoopSet, lower
 
 struct Ex{T, Tup} end
 
-function to_type(ex::Expr)
-    Ex{ex.head, Tuple{to_type.(ex.args)...}}
+function to_type(@nospecialize(ex))
+    if ex isa Expr
+        Ex{ex.head, Tuple{to_type.(ex.args)...}}
+    elseif ex isa LineNumberNode
+        nothing
+    else
+        ex
+    end
 end
 
-to_type(x) = x
-to_type(::LineNumberNode) = nothing
+# to_type(x) = x
+# to_type(::LineNumberNode) = nothing
 
 #----------------------------------------------------------------------------------------------------
 
@@ -19,42 +25,39 @@ 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
+function find_vars!(@nospecialize(ex), vars::Set{Symbol}, ivars::Vector{Symbol})
+    if ex isa Expr
+        if ex.head == :(=) && ex.args[1] isa Symbol
+            push!(ivars, ex.args[1])
+        elseif ex.head == :call
+            push!(ivars, ex.args[1])
+        end
+        ex
+    elseif ex isa Symbol && (ex ∉ vars) && (ex ∉ ivars)
+        push!(vars, ex)
+        ex
     else
-        x
+        ex
     end    
 end
 
-find_vars_and_gensym!(x, vars::Set{Symbol}, ivars::Vector{Symbol}) = x
-
 #----------------------------------------------------------------------------------------------------
 
-nt(keys, vals) = NamedTuple{keys, typeof(vals)}(vals)
+# Not sure whether or not it's better to rely on inlining and const prop. I just like to make things explicit.
+@inline nt(::Val{keys}, vals) where {keys} = NamedTuple{keys, typeof(vals)}(vals)
 
 macro _avx(ex)
     D     = Set{Symbol}()
     ivars = Symbol[]
 
-    gex = prewalk(x -> find_vars_and_gensym!(x, D, ivars), ex)
+    gex = prewalk(x -> find_vars!(x, D, ivars), ex)
 
     type_ex = to_type(gex)
 
     tvars  = Tuple(D)
 
     quote
-        kwargs = LoopVectorization.nt($(QuoteNode(tvars)), $(Expr(:tuple, tvars...)))
+        kwargs = LoopVectorization.nt(Val{$(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

src/constructors.jl

@@ -48,6 +48,15 @@ function LoopSet(q::Expr)
     ls
 end
 
+function LoopSet(q::Expr, types::Dict{Symbol,DataType})
+    q = SIMDPirates.contract_pass(q)
+    ls = LoopSet()
+    copyto!(ls, q, types)
+    resize!(ls.loop_order, num_loops(ls))
+    ls
+end
+
+
 """
     @avx