Check for 0 loopdeps before trying to find array in case of CartesianIndices, fixes #256. Place masks in correct order when manually unroling collections, fixes #257. Add more canonicalization/indexing cleanup.

Author: chriselrod

src/codegen/loopstartstopmanager.jl

@@ -230,22 +230,27 @@ function normalize_offsets!(
   end
   return Int(minoffset)
 end
-function isloopvalue(ls::LoopSet, ind::Symbol)
-  for op ∈ operations(ls)
+function isloopvalue(ls::LoopSet, ind::Symbol, isrooted::Union{Nothing,Vector{Bool}} = nothing)
+  for (i,op) ∈ enumerate(operations(ls))
+    if (isrooted ≢ nothing)
+      isrooted[i] || continue
+    end
     iscompute(op) || continue
     for opp ∈ parents(op)# this is to confirm `ind` still has children
-      (isloopvalue(opp) && instruction(opp).instr === ind) && return true
+      # (isloopvalue(opp) && instruction(opp).instr === ind) && return true
+      if (isloopvalue(opp) && instruction(opp).instr === ind)
+        return true
+      end
     end
   end
   return false
 end
 function cse_constant_offsets!(
-  ls::LoopSet, allarrayrefs::Vector{ArrayReferenceMeta}, allarrayrefsind::Int, name_to_array_map::Vector{Vector{Int}},
-  arrayref_to_name_op_collection::Vector{Vector{Tuple{Int,Int,Int}}}, shouldindbyind::Vector{Bool}
+  ls::LoopSet, allarrayrefs::Vector{ArrayReferenceMeta}, allarrayrefsind::Int, name_to_array_map::Vector{Vector{Int}}, arrayref_to_name_op_collection::Vector{Vector{Tuple{Int,Int,Int}}}
 )
   ar = allarrayrefs[allarrayrefsind]
   # @show ar
-  vptrar = vptr(ar)
+  # vptrar = vptr(ar)
   arrayref_to_name_op = arrayref_to_name_op_collection[allarrayrefsind]
   array_refs_with_same_name = name_to_array_map[first(first(arrayref_to_name_op))]
   us = ls.unrollspecification
@@ -254,7 +259,7 @@ function cse_constant_offsets!(
   strides = getstrides(ar)
   offset = first(indices) === DISCONTIGUOUS
   # gespindoffsets = fill(Symbol(""), length(li))
-  gespinds = Expr(:tuple)
+  gespindsummary = Vector{Tuple{Symbol,Int}}(undef, length(li))
   for i ∈ eachindex(li)
     gespsymbol::Symbol = Symbol("")    
     ii = i + offset
@@ -372,9 +377,10 @@ function cse_constant_offsets!(
       end
     end
     constoffset = normalize_offsets!(ls, i, allarrayrefs, array_refs_with_same_name, arrayref_to_name_op_collection)
-    pushgespind!(gespinds, ls, gespsymbol, constoffset, ind, li, i, check_shouldindbyind(ls, ind, shouldindbyind), true)
+    gespindsummary[i] = (gespsymbol, constoffset)
+    # pushgespind!(gespinds, ls, gespsymbol, constoffset, ind, li, i, check_shouldindbyind(ls, ind, shouldindbyind), true)
   end
-  return gespinds
+  return gespindsummary
 end
 @inline similardims(_, i) = i
 @inline similardims(::CartesianIndices{N}, i) where {N} = VectorizationBase.CartesianVIndex(ntuple(_ -> i, Val{N}()))
@@ -393,10 +399,22 @@ end
 #   end
 #   return nothing
 # end
+function calcgespinds(ls::LoopSet, ar::ArrayReferenceMeta, gespindsummary::Vector{Tuple{Symbol,Int}}, shouldindbyind::Vector{Bool})
+  gespinds = Expr(:tuple)
+  li = ar.loopedindex
+  indices = getindicesonly(ar)
+  for i ∈ eachindex(li)
+    ind = indices[i]
+    gespsymbol, constoffset = gespindsummary[i]
+    pushgespind!(gespinds, ls, gespsymbol, constoffset, ind, li[i], check_shouldindbyind(ls, ind, shouldindbyind), true)
+  end
+  gespinds
+end
+
 function pushgespind!(
-  gespinds::Expr, ls::LoopSet, gespsymbol::Symbol, constoffset::Int, ind::Symbol, li::Vector{Bool}, i::Int, index_by_index::Bool, fromgsp::Bool
+  gespinds::Expr, ls::LoopSet, gespsymbol::Symbol, constoffset::Int, ind::Symbol, isli::Bool, index_by_index::Bool, fromgsp::Bool
 )
-  if li[i]
+  if isli
     if ind === CONSTANTZEROINDEX
       if gespsymbol === Symbol("")
         push!(gespinds.args, staticexpr(constoffset))
@@ -448,13 +466,21 @@ function pushgespind!(
   elseif fromgsp # from gsp means that a loop could be a CartesianIndices, so we may need to expand
     #TODO: broadcast dimensions in case of cartesian indices
     rangesym = ind
+    foundind = false
     for op ∈ operations(ls)
       if name(op) === ind
-        loopsym = first(loopdependencies(op))
-        rangesym = getloop(ls, loopsym).rangesym
+        loopdeps = loopdependencies(op)
+        foundind = true
+        if length(loopdeps) ≠ 0
+          rangesym = getloop(ls, first(loopdeps)).rangesym
+        else
+          isconstantop(op) || throw(LoopError("Please file an issue with LoopVectorization.jl with a reproducer; tried to eliminate a non-constant operation."))
+          rangesym = name(op)
+        end
+        break
       end
     end
-    @assert rangesym ≢ ind
+    @assert foundind
     if rangesym === Symbol("") # there is no rangesym, must be statically sized.
       pushgespsym!(gespinds, gespsymbol, constoffset)
     else
@@ -518,32 +544,32 @@ function use_loop_induct_var!(
   vptrar = vptr(ar)
   # @show ar
   Wisz = false#ls.vector_width == 0
-  for i ∈ eachindex(li)
+  for (i,isli) ∈ enumerate(li)
     ii = i + offset
     ind = indices[ii]
     Wisz && push!(gespinds.args, staticexpr(0)) # wrong for `@_avx`...
     if !li[i] # if it wasn't set
       uliv[i] = 0
       push!(offsetprecalc_descript.args, 0)
-      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, li, i, true, false)
+      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, isli, true, false)
     elseif ind === CONSTANTZEROINDEX
       uliv[i] = 0
       push!(offsetprecalc_descript.args, 0)
-      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, li, i, true, false)
+      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, isli, true, false)
     elseif isbroadcast ||
       ((isone(ii) && (last(looporder) === ind)) && !(otherindexunrolled(ls, ind, ar)) ||
       multiple_with_name(vptrar, allarrayrefs)) ||
       (iszero(ls.vector_width) && isstaticloop(getloop(ls, ind)))# ||
       # Not doing normal offset indexing
       uliv[i] = -findfirst(Base.Fix2(===,ind), looporder)::Int
       push!(offsetprecalc_descript.args, 0) # not doing offset indexing, so push 0
-      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, li, i, true, false)
+      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, isli, true, false)
     else
       uliv[i] = findfirst(Base.Fix2(===,ind), looporder)::Int
       loop = getloop(ls, ind)
       push!(offsetprecalc_descript.args, max(5,us.u₁+1,us.u₂+1))
       use_offsetprecalc = true
-      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, li, i, false, false)
+      Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, ind, isli, false, false)
     end
     # cases for pushgespind! and loopval!
     # if !isloopval, same as before

src/codegen/lower_load.jl

@@ -477,7 +477,7 @@ function lower_load_collection!(
                 uinds = Expr(:call, unrollcurl₂, inds)
                 loadexpr = copy(loadexpr)
                 loadexpr.args[3] = Expr(:call, unrollcurl₂, inds)
-                (((u+1) == u₁) & masklast) && push!(loadexpr.args, MASKSYMBOL)
+                (((u+1) == u₁) & masklast) && insert!(loadexpr.args, length(loadexpr.args)-1, MASKSYMBOL) # 1 for `falseexpr` pushed at end
             end
             # unpack_collection!(q, ls, opidmap, idsformap, ua, loadexpr, collectionname, op, false)
             push!(q.args, Expr(:(=), collectionname_u, Expr(:call, gf, loadexpr, 1)))

src/codegen/lower_store.jl

@@ -122,7 +122,7 @@ function lower_store_collection!(
             lastiter = (u+1) == u₁
             storeexpr_tmp = if lastiter
                 storeexpr
-                (((u+1) == u₁) & masklast) && push!(storeexpr.args, MASKSYMBOL)
+                (((u+1) == u₁) & masklast) && insert!(storeexpr.args, length(storeexpr.args)-3, MASKSYMBOL) # 3 for falseexpr, aliasexpr, falseexpr
                 storeexpr
             else
                 copy(storeexpr)

src/condense_loopset.jl

@@ -159,6 +159,7 @@ end
 @inline zerorangestart(r::AbstractUnitRange) = Zero():One():(maybestaticlast(r)-maybestaticfirst(r))
 @inline zerorangestart(r::AbstractRange) = Zero():static_step(r):(maybestaticlast(r)-maybestaticfirst(r))
 @inline zerorangestart(r::CartesianIndices) = CartesianIndices(map(zerorangestart, r.indices))
+@inline zerorangestart(r::ArrayInterface.OptionallyStaticUnitRange{StaticInt{1}}) = CloseOpen(last(r)) 
 
 function loop_boundary!(q::Expr, ls::LoopSet, loop::Loop, shouldindbyind::Bool)
   if isstaticloop(loop) || loop.rangesym === Symbol("")
@@ -309,12 +310,12 @@ function findfirstcontaining(ref, ind)
   end
   0
 end
-function should_zerorangestart(ls::LoopSet, allarrayrefs::Vector{ArrayReferenceMeta}, name_to_array_map::Vector{Vector{Int}})
+function should_zerorangestart(ls::LoopSet, allarrayrefs::Vector{ArrayReferenceMeta}, name_to_array_map::Vector{Vector{Int}}, isrooted::Vector{Bool})
   loops = ls.loops
   shouldindbyind = fill(false, length(loops))
   for (i,loop) ∈ enumerate(loops)
     ind = loop.itersymbol
-    if isloopvalue(ls, ind)
+    if isloopvalue(ls, ind, isrooted)
       # we don't zero the range if it is used as a loopvalue
       shouldindbyind[i] = true
       continue
@@ -355,10 +356,11 @@ end
 # 2) decide whether to gesp that loopstart inside `add_grouped_strided_pointer`
 function add_grouped_strided_pointer!(extra_args::Expr, ls::LoopSet)
   allarrayrefs, name_to_array_map, unique_to_name_and_op_map = uniquearrayrefs_csesummary(ls)
-  shouldindbyind = should_zerorangestart(ls, allarrayrefs, name_to_array_map)
   # @show allarrayrefs
   gsp = Expr(:call, lv(:grouped_strided_pointer))
   tgarrays = Expr(:tuple)
+  # refs_to_gesp = ArrayReferenceMeta[]
+  gespsummaries = Tuple{Int,Vector{Tuple{Symbol,Int}}}[]
   i = 0
   preserve_assignment = Expr(:tuple); preserve = Symbol[];
   @unpack equalarraydims, refs_aliasing_syms = ls
@@ -383,17 +385,24 @@ function add_grouped_strided_pointer!(extra_args::Expr, ls::LoopSet)
     duplicate && continue
     duplicate_map[j] = (i += 1)
     found = false
-    for j ∈ eachindex(allarrayrefs)
-      if sameref(allarrayrefs[j], ref)
-        gespinds = cse_constant_offsets!(ls, allarrayrefs, j, name_to_array_map, unique_to_name_and_op_map, shouldindbyind)
-        push!(tgarrays.args, Expr(:call, lv(:gespf1), vpref, gespinds))
+    for k ∈ eachindex(allarrayrefs)
+      if sameref(allarrayrefs[k], ref)
+        gespindsummary = cse_constant_offsets!(ls, allarrayrefs, k, name_to_array_map, unique_to_name_and_op_map)
+        push!(gespsummaries, (k, gespindsummary))
         found = true
         break
       end
     end
     @assert found
     push!(preserve, presbufsym(ref.ref.array))
   end
+  roots = getroots(ls)
+  shouldindbyind = should_zerorangestart(ls, allarrayrefs, name_to_array_map, roots)
+  for (k,gespindsummary) ∈ gespsummaries
+    ref = allarrayrefs[k]
+    gespinds = calcgespinds(ls, ref, gespindsummary, shouldindbyind)
+    push!(tgarrays.args, Expr(:call, lv(:gespf1), vptr(ref), gespinds))
+  end
   push!(gsp.args, tgarrays)
   matcheddims = Expr(:tuple)
   for (vptrs,dims) ∈ equalarraydims
@@ -409,7 +418,7 @@ function add_grouped_strided_pointer!(extra_args::Expr, ls::LoopSet)
   gsps = gensym!(ls, "#grouped#strided#pointer#")
   push!(extra_args.args, gsps)
   pushpreamble!(ls, Expr(:(=), gsps, Expr(:call, GlobalRef(Core,:getfield), gsp, 1)))
-  preserve, shouldindbyind
+  preserve, shouldindbyind, roots
 end
 
 # first_cache() = ifelse(gt(num_cache_levels(), StaticInt{2}()), StaticInt{2}(), StaticInt{1}())
@@ -468,12 +477,11 @@ end
 # Try to condense in type stable manner
 function generate_call(ls::LoopSet, (inline,u₁,u₂)::Tuple{Bool,Int8,Int8}, thread::UInt, debug::Bool = false)
     extra_args = Expr(:tuple)
-    preserve, shouldindbyind = add_grouped_strided_pointer!(extra_args, ls)
+    preserve, shouldindbyind, roots = add_grouped_strided_pointer!(extra_args, ls)
 
     operation_descriptions = Expr(:tuple)
     varnames = Symbol[]; ids = Vector{Int}(undef, length(operations(ls)))
     ops = operations(ls)
-    roots = getroots(ls)
     length(ls.includedactualarrays) == 0 || remove_outer_reducts!(roots, ls)
     for op ∈ ops
         instr::Instruction = instruction(op)

src/modeling/graphs.jl

@@ -746,15 +746,35 @@ end
 
 operations(ls::LoopSet) = ls.operations
 function pushop!(ls::LoopSet, op::Operation, var::Symbol = name(op))
-    for opp ∈ operations(ls)
-        if matches(op, opp)
-            ls.opdict[var] = opp
-            return opp
-        end
+  if iscompute(op) && length(loopdependencies(op)) == 0
+    op.node_type = constant
+    opdef = callexpr(instruction(op))
+    opparents = parents(op)
+    mangledname = Symbol('#', instruction(op).instr, '#')
+    while length(opparents) > 0
+      oppname = name(popfirst!(opparents))
+      mangledname = Symbol(mangledname, oppname, '#')
+      push!(opdef.args, oppname)
+      # if opp.instruction == LOOPCONSTANT 
+      #   push!(opdef.args, name(opp))
+      # else
+
+      # end
+    end
+    op.mangledvariable = mangledname
+    pushpreamble!(ls, Expr(:(=), name(op), opdef))
+    op.instruction = LOOPCONSTANT
+    push!(ls.preamble_symsym, (identifier(op), name(op)))
+  end
+  for opp ∈ operations(ls)
+    if matches(op, opp)
+      ls.opdict[var] = opp
+      return opp
     end
-    push!(ls.operations, op)
-    ls.opdict[var] = op
-    op
+  end
+  push!(ls.operations, op)
+  ls.opdict[var] = op
+  op
 end
 function add_block!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
     for x ∈ ex.args

src/parse/add_constants.jl

@@ -16,8 +16,6 @@ function add_constant!(ls::LoopSet, var::Number, elementbytes::Int = 8)
     op = Operation(length(operations(ls)), gensym!(ls, "loopconstnumber"), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
     ops = operations(ls)
     typ = var isa Integer ? HardInt : HardFloat
-    rop = pushop!(ls, op)
-    rop === op || return rop
     if iszero(var)
         for (id,typ_) ∈ ls.preamble_zeros
             (instruction(ops[id]) == LOOPCONSTANT && typ == typ_) && return ops[id]
@@ -37,6 +35,9 @@ function add_constant!(ls::LoopSet, var::Number, elementbytes::Int = 8)
         end
         push!(ls.preamble_symfloat, (identifier(op), var))
     end
+    rop = pushop!(ls, op)
+    rop === op || return rop
+    pushpreamble!(ls, Expr(:(=), name(op), var))
     rop
 end
 function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementbytes::Int)