Fix test failure

Author: chriselrod

src/codegen/lower_load.jl

@@ -66,7 +66,7 @@ function add_prefetches!(q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, pr
         prefetchstride *= gethint(prefetchloop_step)
     end
     offsets[prefetchind] = inner_offset + prefetchstride
-    gespinds = mem_offset_u(op, td, indices_calculated_by_pointer_offsets(ls, op.ref), false, 0, ls)
+    gespinds = mem_offset_u(op, td, indices_calculated_by_pointer_offsets(ls, op.ref), false, 0, ls, false)
     offsets[prefetchind] = inner_offset
     ptr = vptr(op)
     gptr = Symbol(ptr, "##GESPEDPREFETCH##")
@@ -151,12 +151,13 @@ function lower_load_no_optranslation!(
     loadexpr = Expr(:call, lv(:_vload), sptr(op), inds)
     add_memory_mask!(loadexpr, op, td, mask, ls, 0)
     push!(loadexpr.args, falseexpr, rs) # unaligned load
+    # @show op loadexpr
     push!(q.args, Expr(:(=), mvar, loadexpr))
   elseif (u₁ > 1) & opu₁
     t = Expr(:tuple)
     sptrsym = sptr!(q, op)
     for u ∈ 1:u₁
-      inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, u-1, ls)
+      inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, u-1, ls, false)
       loadexpr = Expr(:call, lv(:_vload), sptrsym, inds)
       domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
       add_memory_mask!(loadexpr, op, td, domask, ls, u)
@@ -166,7 +167,7 @@ function lower_load_no_optranslation!(
     end
     push!(q.args, Expr(:(=), mvar, Expr(:call, lv(:VecUnroll), t)))
   else
-    inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls)
+    inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls, false#= not unrolled =#)
     loadexpr = Expr(:call, lv(:_vload), sptr(op), inds)
     add_memory_mask!(loadexpr, op, td, mask, ls, 0)
     push!(loadexpr.args, falseexpr, rs)
@@ -209,7 +210,7 @@ function lower_load_for_optranslation!(
     # abs of steps are equal
     equal_steps = (step₁ == step₂) ⊻ (posindicator ≠ 0x03)
     _td = UnrollArgs(u₁loop, u₂loop, vloop, u₁, u₂max, Core.ifelse(equal_steps, 0, u₂max - 1))
-    gespinds = mem_offset(op, _td, inds_by_ptroff, false, ls)
+    gespinds = mem_offset(op, _td, inds_by_ptroff, false, ls, false)
     ptr = vptr(op)
     gptr = Symbol(ptr, "##GESPED##")
     for i ∈ eachindex(gespinds.args)
@@ -457,7 +458,7 @@ function lower_load_collection!(
     # construct dummy unrolled loop
     offset_dummy_loop = Loop(first(opindices), MaybeKnown(1), MaybeKnown(1024), MaybeKnown(1), Symbol(""), Symbol(""))
     unrollcurl₂ = unrolled_curly(op, nouter, offset_dummy_loop, vloop, mask, 1) # interleave always 1 here
-    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls)
+    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls, false)
     falseexpr = Expr(:call, lv(:False)); rs = staticexpr(reg_size(ls));
 
     opu₁, opu₂ = isunrolled_sym(op, u₁loopsym, u₂loopsym, vloopsym, ls)
@@ -499,7 +500,7 @@ function lower_load_collection!(
         for u ∈ 0:u₁-1
             collectionname_u = Symbol(collectionname, :_, u)
             if u ≠ 0
-                inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, u, ls)
+                inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, u, ls, false)
                 uinds = Expr(:call, unrollcurl₂, inds)
                 loadexpr = copy(loadexpr)
                 loadexpr.args[3] = Expr(:call, unrollcurl₂, inds)

src/codegen/lower_memory_common.jl

@@ -139,7 +139,7 @@ unrolled loads are calculated as offsets with respect to an initial gesp. This h
 Therefore, unrolled === true results in inds being ignored.
 _mm means to insert `mm`s.
 """
-function mem_offset(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vector{Bool}, _mm::Bool, ls::LoopSet)
+function mem_offset(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vector{Bool}, _mm::Bool, ls::LoopSet, preserve_vecunroll::Bool)
   # @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
   ret = Expr(:tuple)
   indices = getindicesonly(op)
@@ -163,19 +163,17 @@ function mem_offset(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vect
       # addoffset!(ret, newname, stride, offset, _mmi)
       _mmi = indvectorized && parent !== op && (!isvectorized(parent))
       @assert !_mmi "Please file an issue with an example of how you got this."
-      if isu₁unrolled(parent) & (td.u₁ > 1)
+      if (isu₁unrolled(parent) & (td.u₁ > 1)) & (!preserve_vecunroll)
         gf = GlobalRef(Core,:getfield)
         firstnew = Expr(:call, gf, Expr(:call, gf, newname, 1), 1, false)
         if isvectorized(parent) & (!_mm)
           firstnew = Expr(:call, lv(:unmm), firstnew)
         end
         addoffset!(ret, 0, firstnew, offset, false)
+      elseif isvectorized(parent) & (!_mm)
+        addoffset!(ret, 0, Expr(:call, lv(:unmm), newname), offset, false)
       else
-        if isvectorized(parent) & (!_mm)
-          addoffset!(ret, 0, Expr(:call, lv(:unmm), newname), offset, false)
-        else
-          addoffset!(ret, 0, newname, offset, false)
-        end
+        addoffset!(ret, 0, newname, offset, false)
       end
     end
   end
@@ -275,18 +273,18 @@ function unrolled_curly(op::Operation, u₁::Int, u₁loop::Loop, vloop::Loop, m
   end
 end
 function unrolledindex(op::Operation, td::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool}, ls::LoopSet)
-    @unpack u₁, u₁loopsym, u₁loop, vloop = td
-    isone(u₁) && return mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls)
-    any(==(u₁loopsym), getindicesonly(op)) || return mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls)
+  @unpack u₁, u₁loopsym, u₁loop, vloop = td
+  isone(u₁) && return mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls, false)
+  any(==(u₁loopsym), getindicesonly(op)) || return mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls, true)
 
-    unrollcurl = unrolled_curly(op, u₁, u₁loop, vloop, mask)
-    ind = mem_offset_u(op, td, inds_calc_by_ptr_offset, false, 0, ls)
-    Expr(:call, unrollcurl, ind)
+  unrollcurl = unrolled_curly(op, u₁, u₁loop, vloop, mask)
+  ind = mem_offset_u(op, td, inds_calc_by_ptr_offset, false, 0, ls, false)
+  Expr(:call, unrollcurl, ind)
 end
 
 function mem_offset_u(
-  op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vector{Bool}, _mm::Bool, incr₁::Int, ls::LoopSet
-  )
+  op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vector{Bool}, _mm::Bool, incr₁::Int, ls::LoopSet, preserve_vecunroll::Bool
+)
   @assert accesses_memory(op) "Computing memory offset only makes sense for operations that access memory."
   @unpack u₁loopsym, u₂loopsym, vloopsym, u₁step, u₂step, vstep, suffix = td
 
@@ -299,7 +297,7 @@ function mem_offset_u(
   # allbasezero = all(inds_calc_by_ptr_offset) && all(iszero, offsets)
   loopedindex = op.ref.loopedindex
   if iszero(incr₁) & iszero(incr₂)
-    return mem_offset(op, td, inds_calc_by_ptr_offset, _mm, ls)
+    return mem_offset(op, td, inds_calc_by_ptr_offset, _mm, ls, preserve_vecunroll)
     # append_inds!(ret, indices, loopedindex)
   else
     for (n,ind) ∈ enumerate(indices)
@@ -334,7 +332,7 @@ function mem_offset_u(
             newname_unmm = mulexpr(newname_unmm,stride)
           end
           addoffset!(ret, 0, newname_unmm, offset, false)
-        elseif isu₁unrolled(parent) & (td.u₁ > 1)
+        elseif (isu₁unrolled(parent) & (td.u₁ > 1)) && !preserve_vecunroll
           gf = GlobalRef(Core,:getfield)
           firstnew = Expr(:call, gf, Expr(:call, gf, newname, 1), incr₁+1, false)
           if stride ≠ 1
@@ -420,20 +418,6 @@ function add_memory_mask!(memopexpr::Expr, op::Operation, td::UnrollArgs, mask::
   nothing
 end
 
-# varassignname(var::Symbol, u::Int, isunrolled::Bool) = isunrolled ? Symbol(var, u) : var
-# # name_memoffset only gets called when vectorized
-# function name_memoffset(var::Symbol, op::Operation, td::UnrollArgs, u₁unrolled::Bool, inds_calc_by_ptr_offset::Vector{Bool}, ls::LoopSet)
-#     @unpack u₁, u₁loopsym, u₂loopsym, suffix = td
-#     if (suffix == -1) && u₁ < 0 # u₁ == -1 sentinel value meaning not unrolled
-#         name = var
-#         mo = mem_offset(op, td, inds_calc_by_ptr_offset, true, 0, ls)
-#     else
-#         name = u₁unrolled ? Symbol(var, u₁) : var
-#         mo = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0, ls)
-#     end
-#     name, mo
-# end
-
 function condvarname_and_unroll(cond::Operation, u₁loop::Symbol, u₂loop::Symbol, vloop::Symbol, suffix::Int, opu₂::Bool, ls::LoopSet)
     condvar, condu₁, condu₂ = variable_name_and_unrolled(cond, u₁loop, u₂loop, vloop, Core.ifelse(opu₂, suffix, -1), ls)
     condvar, condu₁

src/codegen/lower_store.jl

@@ -92,7 +92,7 @@ function lower_store_collection!(
 
     offset_dummy_loop = Loop(first(getindices(op)), MaybeKnown(1), MaybeKnown(1024), MaybeKnown(1), Symbol(""), Symbol(""))
     unrollcurl₂ = unrolled_curly(op, nouter, offset_dummy_loop, vloop, mask, 1)
-    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls)
+    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls, false)
     falseexpr = Expr(:call, lv(:False));
     aliasexpr = falseexpr;
     # trueexpr = Expr(:call, lv(:True));
@@ -150,7 +150,7 @@ function lower_store_collection!(
             end
             storeexpr_tmp.args[3] = Expr(:call, lv(:VecUnroll), vut)
             if u ≠ 0
-                storeexpr_tmp.args[4] = Expr(:call, unrollcurl₂, mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, u, ls))
+                storeexpr_tmp.args[4] = Expr(:call, unrollcurl₂, mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, u, ls, false))
             end
             push!(q.args, storeexpr_tmp)
         end
@@ -216,7 +216,7 @@ function lower_store!(
       data_u₁ && push!(q.args, Expr(:(=), mvard, Expr(:call, lv(:data), mvar)))
       sptrsym = sptr!(q, op)
       for u ∈ 1:u₁
-        inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, u-1, ls)
+        inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, u-1, ls, false)
         # @show isu₁unrolled(opp), opp
         storeexpr = if data_u₁
           if reductfunc === Symbol("")
@@ -235,7 +235,7 @@ function lower_store!(
         push!(q.args, storeexpr)
       end
     else
-      inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, 0, ls)
+      inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, 0, ls, false)
       storeexpr = if reductfunc === Symbol("")
         Expr(:call, lv(:_vstore!), sptr(op), mvar, inds)
       else
@@ -300,7 +300,7 @@ function lower_tiled_store!(blockq::Expr, op::Operation, ls::LoopSet, ua::Unroll
         push!(tup.args, Symbol(variable_name(opp, ifelse(isu₂, t, -1)), '_', u))
     end
     vut = Expr(:call, lv(:VecUnroll), tup) # `VecUnroll` of `VecUnroll`s
-    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls)
+    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false, 0, ls, false)
     unrollcurl₂ = unrolled_curly(op, u₂, u₂loop, vloop, mask)
     falseexpr = Expr(:call, lv(:False));
     aliasexpr = falseexpr;

src/parse/add_constants.jl

@@ -109,7 +109,7 @@ function add_constant_vload!(ls::LoopSet, op::Operation, mpref::ArrayReferenceMe
     for ind ∈ getindicesonly(op)
       ensure_constant_lowered!(ls, mpref, ind)
     end
-    push!(vloadcall.args, mem_offset(op, UnrollArgs(dummyloop, dummyloop, dummyloop, 0, 0, 0), fill(false,nindices), true, ls))
+    push!(vloadcall.args, mem_offset(op, UnrollArgs(dummyloop, dummyloop, dummyloop, 0, 0, 0), fill(false,nindices), true, ls, false))
   end
   push!(vloadcall.args, Expr(:call, lv(:False)), staticexpr(reg_size(ls)))
   pushpreamble!(ls, Expr(:(=), temp, vloadcall))