Merge pull request #75 from timholy/teh/names

Preserve names of arrays and iteration variables across generated boundary

Author: chriselrod

src/condense_loopset.jl

@@ -4,11 +4,13 @@
 Base.:|(u::Unsigned, it::IndexType) = u | UInt8(it)
 Base.:(==)(u::Unsigned, it::IndexType) = (u % UInt8) == UInt8(it)
 
-struct ArrayRefStruct
+struct ArrayRefStruct{array,ptr}
     index_types::UInt64
     indices::UInt64
     offsets::UInt64
 end
+array(ar::ArrayRefStruct{a,p}) where {a,p} = a
+ptr(ar::ArrayRefStruct{a,p}) where {a,p}   = p
 
 function findindoradd!(v::Vector{T}, s::T) where {T}
     ind = findfirst(sᵢ -> sᵢ == s, v)
@@ -43,7 +45,7 @@ function ArrayRefStruct(ls::LoopSet, mref::ArrayReferenceMeta, arraysymbolinds::
             end
         end
     end
-    ArrayRefStruct( index_types, indices, offsets )
+    ArrayRefStruct{mref.ref.array,mref.ptr}( index_types, indices, offsets )
 end
 
 struct OperationStruct <: AbstractLoopOperation
@@ -189,14 +191,15 @@ end
 @inline array_wrapper(A::SubArray) = A.indices
 
 
-
+# If you change the number of arguments here, make commensurate changes
+# to the `insert!` locations in `setup_call_noinline`.
 @generated function __avx__!(
-    ::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, lb::LB,
+    ::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, lb::LB,
     ::Val{AR}, ::Val{D}, ::Val{IND}, subsetvals, arraydescript, vargs::Vararg{<:Any,N}
-) where {UT, OPS, ARF, AM, LB, N, AR, D, IND}
+) where {UT, OPS, ARF, AM, LPSYM, LB, N, AR, D, IND}
     num_vptrs = length(ARF.parameters)::Int
     vptrs = [gensym(:vptr) for _ ∈ 1:num_vptrs]
-    call = Expr(:call, lv(:_avx_!), Val{UT}(), OPS, ARF, AM, :lb)
+    call = Expr(:call, lv(:_avx_!), Val{UT}(), OPS, ARF, AM, LPSYM, :lb)
     for n ∈ 1:num_vptrs
         push!(call.args, vptrs[n])
     end
@@ -241,21 +244,22 @@ function generate_call(ls::LoopSet, IUT, debug::Bool = false)
     foreach(ref -> push!(arrayref_descriptions.args, ArrayRefStruct(ls, ref, arraysymbolinds)), ls.refs_aliasing_syms)
     argmeta = argmeta_and_consts_description(ls, arraysymbolinds)
     loop_bounds = loop_boundaries(ls)
+    loop_syms = Expr(:curly, :Tuple, map(QuoteNode, ls.loopsymbols)...)
     inline, U, T = IUT
     if inline | debug
         func = debug ? lv(:_avx_loopset_debug) : lv(:_avx_!)
         lbarg = debug ? Expr(:call, :typeof, loop_bounds) : loop_bounds
         q = Expr(
             :call, func, Expr(:call, Expr(:curly, :Val, (U,T))),
-            operation_descriptions, arrayref_descriptions, argmeta, lbarg
+            operation_descriptions, arrayref_descriptions, argmeta, loop_syms, lbarg
         )
         debug && deleteat!(q.args, 2)
         foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
     else
         arraydescript = Expr(:tuple)
         q = Expr(
             :call, lv(:__avx__!), Expr(:call, Expr(:curly, :Val, (U,T))),
-            operation_descriptions, arrayref_descriptions, argmeta, loop_bounds, arraydescript
+            operation_descriptions, arrayref_descriptions, argmeta, loop_syms, loop_bounds, arraydescript
         )
         for array ∈ ls.includedactualarrays
             push!(q.args, Expr(:call, lv(:unwrap_array), array))
@@ -315,10 +319,10 @@ function setup_call_noinline(ls::LoopSet, U = zero(Int8), T = zero(Int8))
         end
         push!(q.args, ex)
     end
-    insert!(call.args, 7, Expr(:call, Expr(:curly, :Val, vptrarrays)))
-    insert!(call.args, 8, Expr(:call, Expr(:curly, :Val, vptrsubsetdims)))
-    insert!(call.args, 9, Expr(:call, Expr(:curly, :Val, vptrindices)))
-    insert!(call.args, 10, vptrsubsetvals)
+    insert!(call.args, 8, Expr(:call, Expr(:curly, :Val, vptrarrays)))
+    insert!(call.args, 9, Expr(:call, Expr(:curly, :Val, vptrsubsetdims)))
+    insert!(call.args, 10, Expr(:call, Expr(:curly, :Val, vptrindices)))
+    insert!(call.args, 11, vptrsubsetvals)
     if hasouterreductions
         outer_reducts = Expr(:local)
         for or ∈ ls.outer_reductions

src/reconstruct_loopset.jl

@@ -1,18 +1,18 @@
-function Loop(ls::LoopSet, l::Int, ::Type{UnitRange{Int}})
-    start = gensym(:loopstart); stop = gensym(:loopstop)
+function Loop(ls::LoopSet, l::Int, sym::Symbol, ::Type{UnitRange{Int}})
+    start = gensym(String(sym)*"_loopstart"); stop = gensym(String(sym)*"_loopstop")
     pushpreamble!(ls, Expr(:(=), start, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), Expr(:ref, :lb, l), QuoteNode(:start)))))
     pushpreamble!(ls, Expr(:(=), stop, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), Expr(:ref, :lb, l), QuoteNode(:stop)))))
-    Loop(gensym(:n), 0, 1024, start, stop, false, false)::Loop
+    Loop(sym, 0, 1024, start, stop, false, false)::Loop
 end
-function Loop(ls::LoopSet, l::Int, ::Type{StaticUpperUnitRange{U}}) where {U}
-    start = gensym(:loopstart)
+function Loop(ls::LoopSet, l::Int, sym::Symbol, ::Type{StaticUpperUnitRange{U}}) where {U}
+    start = gensym(String(sym)*"_loopstart")
     pushpreamble!(ls, Expr(:(=), start, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), Expr(:ref, :lb, l), QuoteNode(:L)))))
-    Loop(gensym(:n), U - 1024, U, start, Symbol(""), false, true)::Loop
+    Loop(sym, U - 1024, U, start, Symbol(""), false, true)::Loop
 end
-function Loop(ls::LoopSet, l::Int, ::Type{StaticLowerUnitRange{L}}) where {L}
-    stop = gensym(:loopstop)
+function Loop(ls::LoopSet, l::Int, sym::Symbol, ::Type{StaticLowerUnitRange{L}}) where {L}
+    stop = gensym(String(sym)*"_loopstop")
     pushpreamble!(ls, Expr(:(=), stop, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__, Symbol(@__FILE__)), Expr(:(.), Expr(:ref, :lb, l), QuoteNode(:U)))))
-    Loop(gensym(:n), L, L + 1024, Symbol(""), stop, true, false)::Loop
+    Loop(sym, L, L + 1024, Symbol(""), stop, true, false)::Loop
 end
 # Is there any likely way to generate such a range?
 # function Loop(ls::LoopSet, l::Int, ::Type{StaticLengthUnitRange{N}}) where {N}
@@ -21,19 +21,19 @@ end
 #     pushpreamble!(ls, Expr(:(=), stop, Expr(:call, :(+), start, N - 1)))
 #     Loop(gensym(:n), 0, N, start, stop, false, false)::Loop
 # end
-function Loop(ls, l, ::Type{StaticUnitRange{L,U}}) where {L,U}
-    Loop(gensym(:n), L, U, Symbol(""), Symbol(""), true, true)::Loop
+function Loop(ls, l, sym::Symbol, ::Type{StaticUnitRange{L,U}}) where {L,U}
+    Loop(sym, L, U, Symbol(""), Symbol(""), true, true)::Loop
 end
 
-function add_loops!(ls::LoopSet, LB)
-    loopsyms = [gensym(:n) for _ ∈ eachindex(LB)]
-    for (i,l) ∈ enumerate(LB)
-        add_loop!(ls, Loop(ls, i, l)::Loop)
+function add_loops!(ls::LoopSet, LPSYM, LB)
+    n = max(length(LPSYM), length(LB))
+    for i = 1:n
+        add_loop!(ls, Loop(ls, i, LPSYM[i], LB[i])::Loop)
     end
 end
+
 function ArrayReferenceMeta(
-    ls::LoopSet, ar::ArrayRefStruct, arraysymbolinds::Vector{Symbol}, opsymbols::Vector{Symbol},
-    array::Symbol, vp::Symbol
+    ls::LoopSet, @nospecialize(ar::ArrayRefStruct), arraysymbolinds::Vector{Symbol}, opsymbols::Vector{Symbol}
 )
     index_types = ar.index_types
     indices = ar.indices
@@ -61,8 +61,8 @@ function ArrayReferenceMeta(
         ni -= 1
     end
     ArrayReferenceMeta(
-        ArrayReference(array, index_vec, offset_vec),
-        loopedindex, vp
+        ArrayReference(array(ar), index_vec, offset_vec),
+        loopedindex, ptr(ar)
     )
 end
 
@@ -105,7 +105,7 @@ end
 function create_mrefs!(ls::LoopSet, arf::Vector{ArrayRefStruct}, as::Vector{Symbol}, os::Vector{Symbol}, vargs)
     mrefs = Vector{ArrayReferenceMeta}(undef, length(arf))
     for i ∈ eachindex(arf)
-        ar = ArrayReferenceMeta(ls, arf[i], as, os, Symbol(""), gensym())
+        ar = ArrayReferenceMeta(ls, arf[i], as, os)
         add_mref!(ls, ar, i, vargs[i])
         mrefs[i] = ar
     end
@@ -222,11 +222,11 @@ function sizeofeltypes(v, num_arrays)::Int
     sizeof(T)
 end
 
-function avx_loopset(instr, ops, arf, AM, LB, vargs)
+function avx_loopset(instr, ops, arf, AM, LPSYM, LB, vargs)
     ls = LoopSet(:LoopVectorization)
     num_arrays = length(arf)
     elementbytes = sizeofeltypes(vargs, num_arrays)
-    add_loops!(ls, LB)
+    add_loops!(ls, LPSYM, LB)
     resize!(ls.loop_order, length(LB))
     arraysymbolinds = process_metadata!(ls, AM, length(arf))
     opsymbols = [gensym(:op) for _ ∈ eachindex(ops)]
@@ -245,22 +245,22 @@ function avx_body(ls, UT)
     q
 end
 
-function _avx_loopset_debug(::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LB}, vargs...) where {UT, OPS, ARF, AM, LB}
-    @show OPS ARF AM LB vargs
-    _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LB.parameters, typeof.(vargs))
+function _avx_loopset_debug(::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, ::Type{LB}, vargs...) where {UT, OPS, ARF, AM, LPSYM, LB}
+    @show OPS ARF AM LPSYM LB vargs
+    _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, typeof.(vargs))
 end
-function _avx_loopset(OPSsv, ARFsv, AMsv, LBsv, vargs) where {UT, OPS, ARF, AM, LB}
+function _avx_loopset(OPSsv, ARFsv, AMsv, LPSYMsv, LBsv, vargs)
     nops = length(OPSsv) ÷ 3
     instr = Instruction[Instruction(OPSsv[3i+1], OPSsv[3i+2]) for i ∈ 0:nops-1]
     ops = OperationStruct[ OPSsv[3i] for i ∈ 1:nops ]
     avx_loopset(
         instr, ops,
         ArrayRefStruct[ARFsv...],
-        AMsv, LBsv, vargs
+        AMsv, LPSYMsv, LBsv, vargs
     )
 end
-@generated function _avx_!(::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, lb::LB, vargs...) where {UT, OPS, ARF, AM, LB}
-    ls = _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LB.parameters, vargs)
+@generated function _avx_!(::Val{UT}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, lb::LB, vargs...) where {UT, OPS, ARF, AM, LPSYM, LB}
+    ls = _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, vargs)
     avx_body(ls, UT)
 end