Make index parsing able to handle more complicated expressions without falling back to operations.

Author: chriselrod

src/codegen/lower_load.jl

@@ -187,7 +187,7 @@ end
 @inline lastunroll(vu::VecUnroll) = last(getfield(vu,:data))
 @inline lastunroll(x) = x
 function lower_load_for_optranslation!(
-    q::Expr, op::Operation, ls::LoopSet, td::UnrollArgs, mask::Bool, translationind::Int
+    q::Expr, op::Operation, posindicator::UInt8, ls::LoopSet, td::UnrollArgs, mask::Bool, translationind::Int
 )
     @unpack u₁loop, u₂loop, vloop, u₁, u₂max, suffix = td
 
@@ -206,10 +206,9 @@ function lower_load_for_optranslation!(
     step₁ = gethint(step(u₁loop))
     step₂ = gethint(step(u₂loop))
 
-
     # abs of steps are equal
     #
-    equal_steps = step₁ == step₂
+    equal_steps = (step₁ == step₂) ⊻ (posindicator ≠ 0x03)
     _td = UnrollArgs(u₁loop, u₂loop, vloop, total_unroll, u₂max, Core.ifelse(equal_steps, 0, u₂max - 1))
     gespinds = mem_offset(op, _td, inds_by_ptroff, false)
     ptr = vptr(op)
@@ -286,12 +285,11 @@ function lower_load!(
     if (suffix != -1) && ls.loadelimination[]
         if (u₁ > 1) & (u₂max > 1)
             istr, ispl = isoptranslation(ls, op, UnrollSymbols(u₁loopsym, u₂loopsym, vloopsym))
-        else
-            istr, ispl = 0, false
+            if istr ≠ 0
+                return lower_load_for_optranslation!(q, op, ispl, ls, td, mask, istr)
+            end
         end
-        if !iszero(istr) & ispl
-            return lower_load_for_optranslation!(q, op, ls, td, mask, istr)
-        elseif (suffix > 0) && (u₂loopsym !== vloopsym)
+        if (suffix > 0) && (u₂loopsym !== vloopsym)
             mno, id = maxnegativeoffset(ls, op, u₂loopsym)
             if -suffix < mno < 0 # already checked that `suffix != -1` above
                 varnew = variable_name(op, suffix)

src/codegen/lower_memory_common.jl

@@ -10,14 +10,7 @@ function symbolind(ind::Symbol, op::Operation, td::UnrollArgs)
     @unpack u₁, u₁loopsym, u₂loopsym, u₂max, suffix = td
     parent = parents(op)[id]
     pvar, u₁op, u₂op = variable_name_and_unrolled(parent, u₁loopsym, u₂loopsym, u₂max, suffix)
-    # pvar = if u₂loopsym ∈ loopdependencies(parent)
-    #     variable_name(parent, suffix)
-    # else
-    #     mangledvar(parent)
-    # end
-    u = u₁op ? u₁ : 1
-    ex = Symbol(pvar, '_', u)
-    Expr(:call, lv(:staticm1), ex), parent
+    Symbol(pvar, '_', Core.ifelse(u₁op, u₁, 1)), parent
 end
 
 staticexpr(x::Int) = Expr(:call, Expr(:curly, lv(:Static), x))
@@ -146,6 +139,7 @@ function mem_offset(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vect
         if loopedindex[n]
             addoffset!(ret, indvectorized, vstep, stride, ind, offset, inds_calc_by_ptr_offset[n] | (ind === CONSTANTZEROINDEX))
         else
+            offset -= 1
             newname, parent = symbolind(ind, op, td)
             # _mmi = indvectorized && parent !== op && (!isvectorized(parent))
             # addoffset!(ret, newname, stride, offset, _mmi)
@@ -213,10 +207,10 @@ function unrolled_curly(op::Operation, u₁::Int, u₁loop::Loop, vloop::Loop, m
                 # Expr(:call, Expr(:curly, lv(:Unroll), AU, 1, u₁, AV, intvecsym, M, 1), ind)
                 Expr(:curly, lv(:Unroll), AU, gethint(step(u₁loop)), u₁, AV, intvecsym, M, X)
             else
-                if isone(step(u₁loop))
+                if isone(X)
                     Expr(:curly, lv(:Unroll), AU, intvecsym, u₁, AV, intvecsym, M, X)
                 else
-                    unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, step(u₁loop)))))
+                    unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, X))))
                     Expr(:curly, lv(:Unroll), AU, unrollstepexpr, u₁, AV, intvecsym, M, X)
                 end
             end
@@ -263,11 +257,16 @@ function mem_offset_u(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Ve
             elseif loopedindex[n]
                 addoffset!(ret, indvectorized, vstep, stride, ind, offset, ind_by_offset)
             else
+                offset -= 1
                 newname, parent = symbolind(ind, op, td)
-                # _mmi = _mm && indvectorized && parent !== op && (!isvectorized(parent))
+                _mmi = _mm && indvectorized && parent !== op && (!isvectorized(parent))
                 #                              addoffset!(ret, newname, 1, offset, _mmi)
                 @assert !_mmi "Please file an issue with an example of how you got this."
-                addoffset!(ret, newname, 1, offset, false)
+                if stride == 1
+                    addoffset!(ret, 0, newname, offset, false)
+                else
+                    addoffset!(ret, 0, mulexpr(newname,stride), offset, false)
+                end
             end
         end
     end

src/codegen/lowering.jl

@@ -42,14 +42,15 @@ function lower_block(
     for prepost ∈ 1:2
         # !u₁ && !u₂
         lower!(blockq, ops[1,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, true)
+        # isu₁unrolled, isu₂unrolled, after_loop, n
         opsv1 = ops[1,2,prepost,n]
         opsv2 = ops[2,2,prepost,n]
         if length(opsv1) + length(opsv2) > 0
             nstores = 0
             iszero(length(opsv1)) || (nstores += sum(isstore, opsv1))
             iszero(length(opsv2)) || (nstores += sum(isstore, opsv2))
             # if nstores
-            if (length(opsv1) + length(opsv2) == nstores) # all_u₂_ops_store
+            if (length(opsv1) + length(opsv2) == nstores) && u₂ > 1 # all_u₂_ops_store
                 lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, true) # for u ∈ 0:u₁-1
                 lower_tiled_store!(blockq, opsv1, opsv2, ls, unrollsyms, u₁, u₂, mask)
             else

src/modeling/determinestrategy.jl

@@ -674,17 +674,19 @@ function isoptranslation(ls::LoopSet, op::Operation, unrollsyms::UnrollSymbols)
 
     istranslation = 0
     inds = getindices(op); li = op.ref.loopedindex
-    translationplus = false
     for i ∈ eachindex(li)
         if !li[i]
             opp = findparent(ls, inds[i + (first(inds) === DISCONTIGUOUS)])
-            if instruction(opp).instr ∈ (:+, :-) && isu₁unrolled(opp) && isu₂unrolled(opp)
-                istranslation = i
-                translationplus = instruction(opp).instr === :+
+            if isu₁unrolled(opp) && isu₂unrolled(opp)
+                isadd = instruction(opp).instr === :(+)
+                issub = instruction(opp).instr === :(-)
+                if isadd | issub
+                    return  i, isadd
+                end
             end
         end
     end
-    istranslation, translationplus
+    0, false
 end
 function maxnegativeoffset(ls::LoopSet, op::Operation, u::Symbol)
     mno = typemin(Int)