Support using loop variables for things other than indexing. Fixes #22, fixes #36.

Author: chriselrod

src/add_compute.jl

@@ -156,6 +156,13 @@ function add_compute!(
             else
                 add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes)
             end
+        elseif arg ∈ ls.loopsymbols
+            loopsym = gensym(arg)
+            pushpreamble!(ls, Expr(:(=), loopsym, LoopValue()))
+            loopsymop = add_simple_load!(ls, gensym(loopsym), ArrayReference(loopsym, [arg]), elementbytes)
+            push!(ls.syms_aliasing_refs, name(loopsymop))
+            push!(ls.refs_aliasing_syms, loopsymop.ref)
+            pushparent!(parents, deps, reduceddeps, loopsymop)
         else
             add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes)
         end

src/add_constants.jl

@@ -12,7 +12,7 @@ function add_constant!(ls::LoopSet, var::Symbol, mpref::ArrayReferenceMetaPositi
     op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS, mpref.mref)
     add_vptr!(ls, op)
     temp = gensym(:intermediateconstref)
-    pushpreamble!(ls, Expr(:(=), temp, Expr(:call, lv(:load), mpref.mref.ptr, mem_offset(op, UnrollArgs(zero(Int32), Symbol(""), Symbol(""), nothing)))))
+    pushpreamble!(ls, Expr(:(=), temp, Expr(:call, lv(:load), mpref.mref.ptr, mem_offset(op, UnrollArgs(0, Symbol(""), Symbol(""), nothing)))))
     pushpreamble!(ls, op, temp)
     pushop!(ls, op, temp)
 end

src/add_ifelse.jl

@@ -19,8 +19,7 @@ function add_if!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int = 8, mpr
     add_compute!(ls, LHS, :vifelse, [condop, trueop, falseop], elementbytes)
 end
 
-function add_andblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
-    rhsop = add_compute!(ls, gensym(:iftruerhs), RHS, elementbytes)
+function add_andblock!(ls::LoopSet, condop::Operation, LHS, rhsop::Operation, elementbytes::Int)
     if LHS isa Symbol
         altop = getop(ls, LHS)
         return add_compute!(ls, LHS, :vifelse, [condop, rhsop, altop], elementbytes)
@@ -30,6 +29,14 @@ function add_andblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::I
         throw("Don't know how to assign onto $LHS.")
     end        
 end
+function add_andblock!(ls::LoopSet, condop::Operation, LHS, RHS::Expr, elementbytes::Int)
+    rhsop = add_compute!(ls, gensym(:iftruerhs), RHS, elementbytes)
+    add_andblock!(ls, condop, LHS, rhsop, elementbytes)
+end
+function add_andblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
+    rhsop = getop(ls, RHS)
+    add_andblock!(ls, condop, LHS, rhsop, elementbytes)
+end
 function add_andblock!(ls::LoopSet, condexpr::Expr, condeval::Expr, elementbytes::Int)
     condop = add_compute!(ls, gensym(:mask), condexpr, elementbytes)
     @assert condeval.head === :(=)
@@ -42,8 +49,7 @@ function add_andblock!(ls::LoopSet, ex::Expr, elementbytes::Int)
     add_andblock!(ls, first(ex.args)::Expr, last(ex.args)::Expr, elementbytes)
 end
 
-function add_orblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
-    rhsop = add_compute!(ls, gensym(:iffalserhs), RHS, elementbytes)
+function add_orblock!(ls::LoopSet, condop::Operation, LHS, rhsop::Operation, elementbytes::Int)
     if LHS isa Symbol
         altop = getop(ls, LHS)
         return add_compute!(ls, LHS, :vifelse, [condop, altop, rhsop], elementbytes)
@@ -54,6 +60,14 @@ function add_orblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::In
         throw("Don't know how to assign onto $LHS.")
     end
 end
+function add_orblock!(ls::LoopSet, condop::Operation, LHS, RHS::Expr, elementbytes::Int)
+    rhsop = add_compute!(ls, gensym(:iffalserhs), RHS, elementbytes)
+    add_orblock!(ls, condop, LHS, rhsop, elementbytes)
+end
+function add_orblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
+    rhsop = getop(ls, RHS)
+    add_orblock!(ls, condop, LHS, rhsop, elementbytes)
+end
 function add_orblock!(ls::LoopSet, condexpr::Expr, condeval::Expr, elementbytes::Int)
     condop = add_compute!(ls, gensym(:mask), condexpr, elementbytes)
     @assert condeval.head === :(=)

src/add_loads.jl

@@ -50,3 +50,11 @@ function add_load_getindex!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::In
     add_load!(ls, var, array, rawindices, elementbytes)
 end
 
+
+struct LoopValue end
+@inline VectorizationBase.stridedpointer(::LoopValue) = LoopValue()
+@inline SIMDPirates.vload(::LoopValue, i::Tuple{_MM{W}}) where {W} = SVec(SIMDPirates.vrangeincr(Val{W}(), @inbounds(i[1].i), Val{1}()))
+@inline VectorizationBase.load(::LoopValue, i::Integer) = i + one(i)
+@inline VectorizationBase.load(::LoopValue, i::Tuple{I}) where {I<:Integer} = @inbounds(i[1]) + one(I)
+@inline Base.eltype(::LoopValue) = Int8
+

src/add_stores.jl

@@ -77,12 +77,10 @@ end
 # For now, it is illegal to load from a conditional store.
 # if you want that sort of behavior, do a conditional reassignment, and store that result unconditionally.
 function add_conditional_store!(ls::LoopSet, LHS, condop::Operation, storeop::Operation, elementbytes::Int)
-    array, raw_indices = ref_from_ref(LHS)
-    ref = ArrayReference(array, raw_indices)
-    mref = ArrayReferenceMeta(
-        ref, fill(true, length(getindices(ref)))
-    )
-    ldref = convert(Vector{Symbol}, getindices(ref))
+    array, rawindices = ref_from_ref(LHS)
+    mpref = array_reference_meta!(ls, array, rawindices, elementbytes)
+    mref = mpref.mref
+    ldref = mpref.loopdependencies
 
     pvar = storeop.variable
     id = length(ls.operations)

src/condense_loopset.jl

@@ -75,7 +75,7 @@ function shifted_loopset(ls::LoopSet, loopsyms::Vector{Symbol})
     ld = zero(UInt64) # leading_zeros(ld) >> 2 yields the number of loopdeps
     for d ∈ loopsyms
         ld <<= 4
-        ld |= getloopid(ls, d)
+        ld |= getloopid(ls, d)::Int
     end
     ld
 end

src/determinestrategy.jl

@@ -233,7 +233,7 @@ function solve_tilesize(X, R, UL, TL)
     end
     # The RR + 1 is a hack to get it to favor Uhigh in more scenarios
     Tl = Tlow
-    while RR < Uhigh*Tl*R[1] + Uhigh*R[2]
+    while RR < Uhigh*Tl*R[1] + Uhigh*R[2] && Tl > 1
         Tl -= 1
     end
     tcost_temp = tile_cost(X, Uhigh, Tl, UL, TL)
@@ -440,6 +440,7 @@ function evaluate_cost_tile(
             reg_pressure[4] += rp
         end
     end
+    sum(reg_pressure) > VectorizationBase.REGISTER_COUNT && return 0, 0, Inf
     # @show order, vectorized cost_vec reg_pressure
     # @show solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure)
     U, T, tcost = solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure, W, vectorized)

src/lower_load.jl

@@ -18,7 +18,7 @@ function lower_load_scalar!(
     ptr = refname(op)
     isunrolled = unrolled ∈ loopdeps
     U = isunrolled ? U : 1
-    for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
+    for u ∈ 0:U-1
         varname = varassignname(var, u, isunrolled)
         td = UnrollArgs(u, unrolled, tiled, suffix)
         push!(q.args, Expr(:(=), varname, Expr(:call, lv(:load), ptr, mem_offset_u(op, td))))
@@ -32,16 +32,16 @@ function lower_load_vectorized!(
     loopdeps = loopdependencies(op)
     @assert vectorized ∈ loopdeps
     if unrolled ∈ loopdeps
-        umin = zero(Int32)
+        umin = 0
         U = U
     else
-        umin = -one(Int32)
+        umin = -1
         U = 0
     end
     # Urange = unrolled ∈ loopdeps ? 0:U-1 : 0
     var = variable_name(op, suffix)
     vecnotunrolled = vectorized !== unrolled
-    for u ∈ umin:Base.unsafe_trunc(Int32,U-1)
+    for u ∈ umin:U-1
         td = UnrollArgs(u, unrolled, tiled, suffix)
         pushvectorload!(q, op, var, td, U, W, mask, vecnotunrolled)
     end

src/lower_memory_common.jl

@@ -1,6 +1,6 @@
 
 struct UnrollArgs{T}
-    u::Int32
+    u::Int
     unrolled::Symbol
     tiled::Symbol
     suffix::T
@@ -13,7 +13,7 @@ function parentind(ind::Symbol, op::Operation)
 end
 function symbolind(ind::Symbol, op::Operation, td::UnrollArgs)
     id = parentind(ind, op)
-    id == -1 && return Expr(:call, :-, ind, one(Int32))
+    id == -1 && return Expr(:call, :-, ind, 1)
     @unpack u, unrolled, tiled, suffix = td
     parent = parents(op)[id]
     pvar = if tiled ∈ loopdependencies(parent)
@@ -22,7 +22,7 @@ function symbolind(ind::Symbol, op::Operation, td::UnrollArgs)
         mangledvar(parent)
     end
     pvar = unrolled ∈ loopdependencies(parent) ? Symbol(pvar, u) : pvar
-    Expr(:call, :-, pvar, one(Int32))
+    Expr(:call, :-, pvar, 1)
 end
 function mem_offset(op::Operation, td::UnrollArgs)
     # @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
@@ -103,7 +103,7 @@ end
 #         Expr(:call, :+, q, incr)
 #     end
 # end
-function varassignname(var::Symbol, u::Int32, isunrolled::Bool)
+function varassignname(var::Symbol, u::Int, isunrolled::Bool)
     isunrolled ? Symbol(var, u) : var
 end
 # name_memoffset only gets called when vectorized

src/lower_store.jl

@@ -55,7 +55,7 @@ function lower_conditionalstore_scalar!(
     end
     condunrolled = unrolled ∈ loopdependencies(cond)
     ptr = refname(op)
-    for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
+    for u ∈ 0:U-1
         varname = varassignname(var, u, isunrolled)
         condvarname = varassignname(condvar, u, condunrolled)
         td = UnrollArgs(u, unrolled, tiled, suffix)
@@ -71,10 +71,10 @@ function lower_conditionalstore_vectorized!(
     @assert unrolled ∈ loopdeps
     var = pvariable_name(op, suffix)
     if isunrolled
-        umin = zero(Int32)
+        umin = 0
         U = U
     else
-        umin = -one(Int32)
+        umin = -1
         U = 0
     end
     ptr = refname(op)
@@ -87,7 +87,7 @@ function lower_conditionalstore_vectorized!(
     end
     # @show parents(op) cond condvar
     condunrolled = unrolled ∈ loopdependencies(cond)
-    for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
+    for u ∈ 0:U-1
         td = UnrollArgs(u, unrolled, tiled, suffix)
         name, mo = name_memoffset(var, op, td, W, vecnotunrolled)
         condvarname = varassignname(condvar, u, condunrolled)
@@ -107,7 +107,7 @@ function lower_store_scalar!(
 )
     var = pvariable_name(op, suffix)
     ptr = refname(op)
-    for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
+    for u ∈ 0:U-1
         varname = varassignname(var, u, isunrolled)
         td = UnrollArgs(u, unrolled, tiled, suffix)
         push!(q.args, Expr(:call, lv(:store!), ptr, varname, mem_offset_u(op, td)))
@@ -122,15 +122,15 @@ function lower_store_vectorized!(
     @assert unrolled ∈ loopdeps
     var = pvariable_name(op, suffix)
     if isunrolled
-        umin = zero(Int32)
+        umin = 0
         U = U
     else
-        umin = -one(Int32)
+        umin = -1
         U = 0
     end
     ptr = refname(op)
     vecnotunrolled = vectorized !== unrolled
-    for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
+    for u ∈ 0:U-1
         td = UnrollArgs(u, unrolled, tiled, suffix)
         name, mo = name_memoffset(var, op, td, W, vecnotunrolled)
         instrcall = Expr(:call, lv(:vstore!), ptr, name, mo)