misc fixes

Author: chriselrod

src/codegen/loopstartstopmanager.jl

@@ -250,9 +250,9 @@ function pointermax_index(
                 else
                     staticexpr(stophint - sub)
                 end
-                stride = getstrides(ar)[i]
+                stride = getstrides(ar)[j]
                 if isknown(incr)
-                    stride *= gethint
+                    stride *= gethint(incr)
                 else
                     _ind = mulexpr(_ind, getsym(incr))
                 end
@@ -353,7 +353,7 @@ function append_pointer_maxes!(
         dim = length(getindicesonly(ar))
         # OFFSETPRECALCDEF = true
         # if OFFSETPRECALCDEF
-        strd = getstrides(ar)[dim]
+        strd = getstrides(ar)[ind]
         for sub ∈ 0:submax-1
             ptrcmp = Expr(:call, lv(:gesp), pointercompbase, offsetindex(dim, ind, (submax - sub)*strd, isvectorized, incr))
             push!(loopstart.args, Expr(:(=), maxsym(vptr_ar, sub), ptrcmp))
@@ -373,6 +373,7 @@ function append_pointer_maxes!(
 end
 function append_pointer_maxes!(loopstart::Expr, ls::LoopSet, ar::ArrayReferenceMeta, n::Int, submax::Int, isvectorized::Bool)
     loop = getloop(ls, n)
+    @assert loop.itersymbol == names(ls)[n]
     start = first(loop)
     stop = last(loop)
     incr = step(loop)

src/codegen/lower_compute.jl

@@ -164,6 +164,7 @@ function ifelselastexpr(hasf::Bool, M::Int, vargtypes, K::Int, S::Int, maskearly
         push!(t.args, :(getfield($(vargs[K]), $m, false)))
     end
     # push!(q.args, :(VecUnroll($t)::VecUnroll{$N,$W,$T,$V}))
+    # push!(q.args, Expr(:call, lv(:VecUnroll), t))
     push!(q.args, :(VecUnroll($t)))
     q
 end
@@ -361,7 +362,7 @@ function lower_compute!(
                 newpname = Symbol(newparentname, '_', u₁)
                 push!(q.args, Expr(:(=), newpname, Symbol(parentname, '_', u₁)))
                 # @show newparentop op instruction(newparentop)
-                reduce_expr!(q, newparentname, instruction(newparentop), u₁, -1)
+                reduce_expr!(q, newparentname, instruction(newparentop), u₁, -1, true)
                 push!(q.args, Expr(:(=), Symbol(newparentname, '_', 1), Symbol(newparentname, "##onevec##")))
             end
         end

src/codegen/lower_load.jl

@@ -124,17 +124,17 @@ function lower_load_no_optranslation!(
         push!(loadexpr.args, falseexpr, rs) # unaligned load
         push!(q.args, Expr(:(=), mvar, loadexpr))
     elseif u₁ > 1
-        # t = Expr(:tuple)
-        # for u ∈ 1:u₁
-        let t = u₁, t = q
+        t = Expr(:tuple)
+        for u ∈ 1:u₁
             inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, u-1)
             loadexpr = Expr(:call, lv(:_vload), vptr(op), inds)
-            add_memory_mask!(loadexpr, op, td, mask & ((u == u₁) | isvectorized(op)))
+            domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
+            add_memory_mask!(loadexpr, op, td, domask)
             push!(loadexpr.args, falseexpr, rs)
-            # push!(t.args, loadexpr)
-            push!(t.args, Expr(:(=), mvar, loadexpr))
+            push!(t.args, loadexpr)
+            # push!(q.args, Expr(:(=), mvar, loadexpr))
         end
-        # push!(q.args, Expr(:(=), mvar, Expr(:call, lv(:VecUnroll), t)))
+        push!(q.args, Expr(:(=), mvar, Expr(:call, lv(:VecUnroll), t)))
     else
         inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0)
         loadexpr = Expr(:call, lv(:_vload), vptr(op), inds)
@@ -185,6 +185,8 @@ end
 @inline firstunroll(x) = x
 @inline lastunroll(vu::VecUnroll) = last(getfield(vu,:data))
 @inline lastunroll(x) = x
+@inline unmm(x) = x
+@inline unmm(x::MM) = getfield(x, :i)
 function lower_load_for_optranslation!(
     q::Expr, op::Operation, posindicator::UInt8, ls::LoopSet, td::UnrollArgs, mask::Bool, translationind::Int
 )
@@ -217,7 +219,7 @@ function lower_load_for_optranslation!(
         if i == translationind
             gespinds.args[i] = Expr(:call, lv(Core.ifelse(equal_steps, :firstunroll, :lastunroll)), gespinds.args[i])
         # else
-        #     gespinds.args[i] = Expr(:call, lv(:data), gespinds.args[i])
+        #     gespinds.args[i] = Expr(:call, lv(:unmm), gespinds.args[i])
         end
     end
     push!(q.args, Expr(:(=), gptr, Expr(:call, lv(:gesp), ptr, gespinds)))
@@ -267,6 +269,7 @@ function lower_load_for_optranslation!(
         broadcasted_data = broadcastedname(variable_name_data)
         push!(q.args, :($broadcasted_data = getfield($(broadcastedname(variable_name_u)), 1)))
     end
+    gf = GlobalRef(Core,:getfield)
     for u₂ ∈ 0:u₂max-1
         variable_name_u₂ = Symbol(variable_name(op, u₂), '_', u₁)
         t = Expr(:tuple)
@@ -279,14 +282,14 @@ function lower_load_for_optranslation!(
             else
                 u - u₂ + u₂max - 1
             end
-            push!(t.args, :(getfield($variable_name_data, $uu)))
+            push!(t.args, :($gf($variable_name_data, $uu)))
             if shouldbroadcast
-                push!(tb.args, :(getfield($broadcasted_data, $uu)))
+                push!(tb.args, :($gf($broadcasted_data, $uu)))
             end
         end
-        push!(q.args, :($variable_name_u₂ = VecUnroll($t)))
+        push!(q.args, Expr(:(=), variable_name_u₂, Expr(:call, lv(:VecUnroll), t)))
         if shouldbroadcast
-            push!(q.args, :($(broadcastedname(variable_name_u₂)) = VecUnroll($tb)))
+            push!(q.args, Expr(:(=), broadcastedname(variable_name_u₂), Expr(:call, lv(:VecUnroll), tb)))
         end
     end
     nothing
@@ -324,8 +327,8 @@ function _lower_load!(
 )
     omop = offsetloadcollection(ls)
     batchid, opind = omop.batchedcollectionmap[identifier(op)]
-    # @show batchid == 0 (!isvectorized(op)) rejectinterleave(op, td.vloop, idsformap)
-    if batchid == 0 || (!isvectorized(op)) || (rejectinterleave(op, td.vloop, omop.batchedcollections[batchid]))
+    # @show batchid == 0 (!isvectorized(op)) rejectinterleave(ls, op, td.vloop, idsformap)
+    if batchid == 0 || (!isvectorized(op)) || (rejectinterleave(ls, op, td.vloop, omop.batchedcollections[batchid]))
         lower_load_no_optranslation!(q, ls, op, td, mask, inds_calc_by_ptr_offset)
     elseif opind == 1# only lower loads once
         # I do not believe it is possible for `opind == 1` to be lowered after an  operation depending on a different opind.
@@ -337,9 +340,26 @@ function _lower_load!(
         lower_load_collection!(q, ls, opidmap, idsformap, td, mask, inds_calc_by_ptr_offset)
     end
 end
-function rejectinterleave(op::Operation, vloop::Loop, idsformap::SubArray{Tuple{Int,Int}, 1, Vector{Tuple{Int,Int}}, Tuple{UnitRange{Int}}, true})
+function addive_loopinductvar_only(op::Operation)
+    isloopvalue(op) && return true
+    iscompute(op) || return false
+    additive_instr = (:add_fast, :(+), :vadd, :identity, :sub_fast, :(-), :vsub)
+    Base.sym_in(instruction(op).instr, additive_instr) || return false
+    return all(addive_loopinductvar_only, parents(op))
+end
+
+function rejectinterleave(ls::LoopSet, op::Operation, vloop::Loop, idsformap::SubArray{Tuple{Int,Int}, 1, Vector{Tuple{Int,Int}}, Tuple{UnitRange{Int}}, true})
     vloopsym = vloop.itersymbol; strd = step(vloop)
     isknown(strd) || return true
+    # TODO: reject if there is a vectorized !loopedindex index
+    for (li,ind) ∈ zip(op.ref.loopedindex,getindicesonly(op))
+        li && continue
+        for indop ∈ operations(ls)
+            if (name(indop) === ind) && isvectorized(indop)
+                addive_loopinductvar_only(op) || return true # so that it is `MM`
+            end
+        end
+    end
     (first(getindices(op)) === vloopsym) && (length(idsformap) ≠ first(getstrides(op)) * gethint(strd))
 end
 function lower_load_collection!(
@@ -375,7 +395,7 @@ function lower_load_collection!(
     # not using `add_memory_mask!(storeexpr, op, ua, mask)` because we checked `isconditionalmemop` earlier in `lower_load_collection!`
     (mask && isvectorized(op)) && push!(loadexpr.args, MASKSYMBOL)
     push!(loadexpr.args, falseexpr, rs)
-    collectionname = Symbol(vp, "##collection##number", first(first(idsformap)), "##size##", nouter, "##u₁##", u₁)
+    collectionname = Symbol(vp, "##collection##number#", opidmap[first(first(idsformap))], "#", suffix, "##size##", nouter, "##u₁##", u₁)
     # getfield to extract data from `VecUnroll` object, so we have a tuple
     push!(q.args, Expr(:(=), collectionname, Expr(:call, :getfield, loadexpr, 1)))
     u = Core.ifelse(isu₁unrolled(op), u₁, 1)

src/codegen/lower_memory_common.jl

@@ -81,7 +81,7 @@ function addoffset!(ret::Expr, indvectorized::Bool, vloopstride, indexstride, in
     end
 end
 
-function addoffset!(ret::Expr, indvectorized::Bool, unrolledsteps, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 8 -> 7 args
+function addvectoroffset!(ret::Expr, indvectorized::Bool, unrolledsteps, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 8 -> 7 args
     # if _iszero(unrolledsteps) # if no steps, pass through; should be unreachable
     #     addoffset!(ret, indvectorized, vloopstride, indexstride, index, offset, calcbypointeroffset)
     # else
@@ -98,22 +98,22 @@ function addoffset!(ret::Expr, indvectorized::Bool, unrolledsteps, vloopstride,
     end
 end
 # unrolledloopstride is a stride multiple on `unrolledsteps`
-function addoffset!(
-    ret::Expr, indvectorized::Bool, unrolledsteps::Int, unrolledloopstride, vloopstride, indexstride::Integer, index, offset::Integer, calcbypointeroffset::Bool
-) # 9 -> (7 or 8) args
+function addvectoroffset!(
+    ret::Expr, mm::Bool, unrolledsteps::Int, unrolledloopstride, vloopstride, indexstride::Integer, index, offset::Integer, calcbypointeroffset::Bool, indvectorized::Bool
+) # 10 -> (7 or 8) args
     if unrolledsteps == 0 # neither unrolledloopstride or indexstride can be 0
-        addoffset!(ret, indvectorized, vloopstride, indexstride, index, offset, calcbypointeroffset) # 7 arg
+        addoffset!(ret, mm, vloopstride, indexstride, index, offset, calcbypointeroffset) # 7 arg
     elseif indvectorized
         unrolledsteps *= indexstride
         if isknown(unrolledloopstride)
-            addoffset!(ret, indvectorized, gethint(unrolledloopstride)*unrolledsteps, vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
+            addvectoroffset!(ret, mm, gethint(unrolledloopstride)*unrolledsteps, vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
         elseif unrolledsteps == 1
-            addoffset!(ret, indvectorized, unrolledloopstride, vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
+            addvectoroffset!(ret, mm, unrolledloopstride, vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
         else
-            addoffset!(ret, indvectorized, mulexpr(unrolledloopstride,unrolledsteps), vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
+            addvectoroffset!(ret, mm, mulexpr(unrolledloopstride,unrolledsteps), vloopstride, indexstride, index, offset, calcbypointeroffset) # 8 arg
         end
     else
-        addoffset!(ret, indvectorized, vloopstride, indexstride, index, offset + unrolledsteps, calcbypointeroffset) # 7 arg
+        addoffset!(ret, mm, vloopstride, indexstride, index, offset + unrolledsteps, calcbypointeroffset) # 7 arg
     end
 end
 
@@ -137,15 +137,28 @@ function mem_offset(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Vect
         stride = strides[n] % Int
         @unpack vstep = td
         if loopedindex[n]
-            addoffset!(ret, indvectorized, vstep, stride, ind, offset, inds_calc_by_ptr_offset[n] | (ind === CONSTANTZEROINDEX))
+            addoffset!(ret, indvectorized, vstep, stride, ind, offset, inds_calc_by_ptr_offset[n] | (ind === CONSTANTZEROINDEX)) # 7 arg
         else
             offset -= 1
             newname, parent = symbolind(ind, op, td)
             # _mmi = indvectorized && parent !== op && (!isvectorized(parent))
             # 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."
-            addoffset!(ret, 0, newname, offset, false)
+            if isu₁unrolled(parent) & (td.u₁ > 1)
+                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)
+            else
+                if isvectorized(parent) & (!_mm)
+                    addoffset!(ret, 0, Expr(:call, lv(:unmm), newname), offset, false)
+                else
+                    addoffset!(ret, 0, newname, offset, false)
+                end
+            end
         end
     end
     ret
@@ -249,20 +262,42 @@ function mem_offset_u(op::Operation, td::UnrollArgs, inds_calc_by_ptr_offset::Ve
             ind_by_offset = inds_calc_by_ptr_offset[n] | (ind === CONSTANTZEROINDEX)
             offset = convert(Int, offsets[n])
             stride = convert(Int, strides[n])
-            indvectorized = _mm & (ind === vloopsym)
+            indvectorized = ind === vloopsym
+            indvectorizedmm = _mm & indvectorized
             if ind === u₁loopsym
-                addoffset!(ret, indvectorized, incr₁, u₁step, vstep, stride, ind, offset, ind_by_offset)
+                addvectoroffset!(ret, indvectorizedmm, incr₁, u₁step, vstep, stride, ind, offset, ind_by_offset, indvectorized) # 9 arg
             elseif ind === u₂loopsym
-                addoffset!(ret, indvectorized, incr₂, u₂step, vstep, stride, ind, offset, ind_by_offset) 
+                # if isstore(op)
+                #     @show indvectorized, ind === vloopsym, u₂loopsym, incr₂
+                # end
+                addvectoroffset!(ret, indvectorizedmm, incr₂, u₂step, vstep, stride, ind, offset, ind_by_offset, indvectorized) # 9 arg
             elseif loopedindex[n]
-                addoffset!(ret, indvectorized, vstep, stride, ind, offset, ind_by_offset)
+                addoffset!(ret, indvectorizedmm, vstep, stride, ind, offset, ind_by_offset) # 7 arg
             else
                 offset -= 1
                 newname, parent = symbolind(ind, op, td)
-                _mmi = _mm && indvectorized && parent !== op && (!isvectorized(parent))
+                _mmi = indvectorizedmm && parent !== op && (!isvectorized(parent))
                 #                              addoffset!(ret, newname, 1, offset, _mmi)
                 @assert !_mmi "Please file an issue with an example of how you got this."
-                if stride == 1
+                if isvectorized(parent) & (!_mm)
+                    if isu₁unrolled(parent) & (td.u₁ > 1)
+                        gf = GlobalRef(Core,:getfield)
+                        newname_unmm = Expr(:call, lv(:unmm), Expr(:call, gf, Expr(:call, gf, newname, 1), incr₁+1, false))
+                    else
+                        newname_unmm = Expr(:call, lv(:unmm), newname)
+                    end
+                    if stride ≠ 1
+                        newname_unmm = mulexpr(newname_unmm,stride)
+                    end
+                    addoffset!(ret, 0, newname_unmm, offset, false)
+                elseif isu₁unrolled(parent) & (td.u₁ > 1)
+                    gf = GlobalRef(Core,:getfield)
+                    firstnew = Expr(:call, gf, Expr(:call, gf, newname, 1), incr₁+1, false)
+                    if stride ≠ 1
+                        firstnew = mulexpr(firstnew,stride)
+                    end
+                    addoffset!(ret, 0, firstnew, offset, false)
+                elseif stride == 1
                     addoffset!(ret, 0, newname, offset, false)
                 else
                     addoffset!(ret, 0, mulexpr(newname,stride), offset, false)
@@ -276,7 +311,7 @@ end
 @inline and_last(a, b) = a & b
 @generated function and_last(v::VecUnroll{N}, m) where {N}
     q = Expr(:block, Expr(:meta,:inline), :(vd = data(v)))
-    t = Expr(:tuple)
+    t = Expr(:call, lv(:promote))
     for n ∈ 1:N
         push!(t.args, :(getfield(vd, $n, false)))
     end

src/codegen/lower_store.jl

@@ -26,14 +26,14 @@ function reduce_expr_u₂(toreduct::Symbol, instr::Instruction, u₂::Int)
     end
     Expr(:call, lv(:reduce_tup), reduce_to_onevecunroll(instr), t)
 end
-function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, u₁::Int, u₂::Int)
+function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, u₁::Int, u₂::Int, isu₁unrolled::Bool)
     if u₂ != -1
         _toreduct = Symbol(toreduct, 0)
         push!(q.args, Expr(:(=), _toreduct, reduce_expr_u₂(toreduct, instr, u₂)))
     else
         _toreduct = Symbol(toreduct, '_', u₁)
     end
-    if u₁ == 1
+    if (u₁ == 1) | (~isu₁unrolled)
         push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), _toreduct))
     else
         push!(q.args, Expr(:(=), Symbol(toreduct, "##onevec##"), Expr(:call, lv(reduction_to_single_vector(instr)), _toreduct)))
@@ -105,7 +105,7 @@ function lower_store!(
 
     omop = offsetloadcollection(ls)
     batchid, opind = omop.batchedcollectionmap[identifier(op)]
-    if ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op, vloop, omop.batchedcollections[batchid]))
+    if ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(ls, op, vloop, omop.batchedcollections[batchid]))
         (opind == 1) && lower_store_collection!(q, ls, op, ua, mask, inds_calc_by_ptr_offset)
         return
     end
@@ -142,7 +142,8 @@ function lower_store!(
             else
                 Expr(:call, lv(:_vstore!), lv(reductfunc), vptr(op), mvaru, inds)
             end
-            add_memory_mask!(storeexpr, op, ua, mask & ((u == u₁) | isvectorized(op)))
+            domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
+            add_memory_mask!(storeexpr, op, ua, domask)# & ((u == u₁) | isvectorized(op)))
             push!(storeexpr.args, falseexpr, trueexpr, falseexpr, rs)
             push!(q.args, storeexpr)
         end
@@ -176,7 +177,7 @@ function donot_tile_store(ls::LoopSet, op::Operation, vloop::Loop, reductfunc::S
 
     omop = offsetloadcollection(ls)
     batchid, opind = omop.batchedcollectionmap[identifier(op)]
-    return ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op, vloop, omop.batchedcollections[batchid]))
+    return ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(ls, op, vloop, omop.batchedcollections[batchid]))
 end
 
 # VectorizationBase implements optimizations for certain grouped stores
@@ -212,7 +213,7 @@ function lower_tiled_store!(blockq::Expr, op::Operation, ls::LoopSet, ua::Unroll
         mvar = Symbol(variable_name(opp, t), '_', u)
         push!(tup.args, mvar)
     end
-    vut = :(VecUnroll($tup)) # `VecUnroll` of `VecUnroll`s
+    vut = Expr(:call, lv(:VecUnroll), tup) # `VecUnroll` of `VecUnroll`s
     inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false)
     unrollcurl₂ = unrolled_curly(op, u₂, u₂loop, vloop, mask)
     falseexpr = Expr(:call, lv(:False)); trueexpr = Expr(:call, lv(:True)); rs = staticexpr(reg_size(ls));

src/codegen/lowering.jl

@@ -540,7 +540,7 @@ function reduce_expr!(q::Expr, ls::LoopSet, U::Int)
         var = name(op)
         mvar = mangledvar(op)
         instr = instruction(op)
-        reduce_expr!(q, mvar, instr, u1f, u2f)
+        reduce_expr!(q, mvar, instr, u1f, u2f, isu₁unrolled(op))
         if !iszero(length(ls.opdict))
             if (isu₁unrolled(op) | isu₂unrolled(op))
                 push!(q.args, Expr(:(=), var, Expr(:call, lv(reduction_scalar_combine(instr)), Symbol(mvar, "##onevec##"), var)))