Keep track of parent module of functions, and test @_avx on a function not defined in LoopVectorization. Fixes #31.

Author: chriselrod

src/LoopVectorization.jl

@@ -13,7 +13,7 @@ using MacroTools: prewalk, postwalk
 
 
 export LowDimArray, stridedpointer, vectorizable,
-    @avx, *ˡ, ∗,
+    @avx, @_avx, *ˡ, _avx_!,
     vmap, vmap!
 
 
@@ -39,8 +39,6 @@ include("condense_loopset.jl")
 include("reconstruct_loopset.jl")
 include("constructors.jl")
 
-export @_avx, _avx, @_avx_, avx_!
-
 # include("precompile.jl")
 # _precompile_()
 

src/add_compute.jl

@@ -86,7 +86,7 @@ function add_reduction_update_parent!(
 )
     parent = getop(ls, var, elementbytes)
     isouterreduction = parent.instruction === LOOPCONSTANT
-    Instr = Instruction(instr)
+    Instr = instruction(ls, instr)
     instrclass = reduction_instruction_class(Instr) # key allows for faster lookups
     # if parent is not an outer reduction...
     if !isouterreduction
@@ -164,7 +164,7 @@ function add_compute!(
     if reduction || search_tree(parents, var)
         add_reduction_update_parent!(parents, deps, reduceddeps, ls, var, instr, reduction, elementbytes)
     else
-        op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, parents)
+        op = Operation(length(operations(ls)), var, elementbytes, instruction(ls,instr), compute, deps, reduceddeps, parents)
         pushop!(ls, op, var)
     end
 end

src/broadcast.jl

@@ -180,12 +180,12 @@ end
 # size of dest determines loops
 # function vmaterialize!(
 @generated function vmaterialize!(
-    dest::AbstractArray{T,N}, bc::BC
-) where {T <: Union{Float32,Float64}, N, BC <: Broadcasted}
+    dest::AbstractArray{T,N}, bc::BC, ::Val{Mod}
+) where {T <: Union{Float32,Float64}, N, BC <: Broadcasted, Mod}
     # we have an N dimensional loop.
     # need to construct the LoopSet
     loopsyms = [gensym(:n) for n ∈ 1:N]
-    ls = LoopSet()
+    ls = LoopSet(Mod)
     sizes = Expr(:tuple)
     for (n,itersym) ∈ enumerate(loopsyms)
         Nsym = gensym(:N)
@@ -204,12 +204,12 @@ end
     # ls
 end
 @generated function vmaterialize!(
-    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::BC
-) where {T <: Union{Float32,Float64}, N, A <: AbstractArray{T,N}, BC <: Broadcasted}
+    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::BC, ::Val{Mod}
+) where {T <: Union{Float32,Float64}, N, A <: AbstractArray{T,N}, BC <: Broadcasted, Mod}
     # we have an N dimensional loop.
     # need to construct the LoopSet
     loopsyms = [gensym(:n) for n ∈ 1:N]
-    ls = LoopSet()
+    ls = LoopSet(Mod)
     pushpreamble!(ls, Expr(:(=), :dest, Expr(:call, :parent, :dest′)))
     sizes = Expr(:tuple)
     for (n,itersym) ∈ enumerate(loopsyms)
@@ -229,7 +229,7 @@ end
     # ls
 end
 
-@inline function vmaterialize(bc::Broadcasted)
+@inline function vmaterialize(bc::Broadcasted, ::Val{Mod}) where {Mod}
     ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
-    vmaterialize!(similar(bc, ElType), bc)
+    vmaterialize!(similar(bc, ElType), bc, Val{Mod}())
 end

src/constructors.jl

@@ -6,7 +6,7 @@ function Base.copyto!(ls::LoopSet, q::Expr)
     add_loop!(ls, q)
 end
 
-function add_ci_call!(q::Expr, f, args, syms, i)
+function add_ci_call!(q::Expr, f, args, syms, i, mod = nothing)
     call = Expr(:call, f)
     for arg ∈ @view(args[2:end])
         if arg isa Core.SSAValue
@@ -15,10 +15,11 @@ function add_ci_call!(q::Expr, f, args, syms, i)
             push!(call.args, arg)
         end
     end
+    mod === nothing || push!(call.args, Expr(:call, Expr(:curly, :Val, QuoteNode(mod))))
     push!(q.args, Expr(:(=), syms[i], call))
 end
 
-function substitute_broadcast(q::Expr)
+function substitute_broadcast(q::Expr, mod::Symbol)
     ci = first(Meta.lower(LoopVectorization, q).args).code
     nargs = length(ci)-1
     ex = Expr(:block,)
@@ -30,32 +31,24 @@ function substitute_broadcast(q::Expr)
         if ciₙ.head === :(=)
             push!(ex.args, Expr(:(=), f, syms[((ciₙargs[2])::Core.SSAValue).id]))
         elseif f === GlobalRef(Base, :materialize!)
-            add_ci_call!(ex, lv(:vmaterialize!), ciₙargs, syms, n)
+            add_ci_call!(ex, lv(:vmaterialize!), ciₙargs, syms, n, mod)
         elseif f === GlobalRef(Base, :materialize)
-            add_ci_call!(ex, lv(:vmaterialize), ciₙargs, syms, n)
+            add_ci_call!(ex, lv(:vmaterialize), ciₙargs, syms, n, mod)
         else
             add_ci_call!(ex, f, ciₙargs, syms, n)
         end
     end
     ex
 end
 
-function LoopSet(q::Expr)
+function LoopSet(q::Expr, mod::Symbol = :LoopVectorization)
     q = SIMDPirates.contract_pass(q)
-    ls = LoopSet()
+    ls = LoopSet(mod)
     copyto!(ls, q)
     resize!(ls.loop_order, num_loops(ls))
     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
@@ -91,10 +84,11 @@ true
 
 """
 macro avx(q)
+    mod = Symbol(__module__)
     q2 = if q.head === :for
-        setup_call(LoopSet(q))
+        setup_call(LoopSet(q, mod))
     else# assume broadcast
-        substitute_broadcast(q)
+        substitute_broadcast(q, mod)
     end
     esc(q2)
 end
@@ -136,24 +130,24 @@ macro avx(arg, q)
     @assert q.head === :for
     @assert arg.head === :(=)
     inline, U, T = check_macro_kwarg(arg)
-    esc(setup_call(LoopSet(q), inline, U, T))
+    esc(setup_call(LoopSet(q, Symbol(__module__)), inline, U, T))
 end
 macro avx(arg1, arg2, q)
     @assert q.head === :for
     inline, U, T = check_macro_kwarg(arg1)
     inline, U, T = check_macro_kwarg(arg2, inline, U, T)
-    esc(setup_call(LoopSet(q), inline, U, T))
+    esc(setup_call(LoopSet(q, Symbol(__module__)), inline, U, T))
 end
 
 
 
 macro _avx(q)
-    esc(lower(LoopSet(q)))
+    esc(lower(LoopSet(q, Symbol(__module__))))
 end
 macro _avx(arg, q)
     @assert q.head === :for
     inline, U, T = check_macro_kwarg(arg)
-    esc(lower(LoopSet(q), U, T))
+    esc(lower(LoopSet(q, Symbol(__module__)), U, T))
 end