Remove runtime dispatches in lower_and_split_loops and lower (using JETTest), remove some unused code, test size on LowDimArrays

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.43"
+version = "0.12.44"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/codegen/lower_compute.jl

@@ -375,7 +375,7 @@ function reduce_parent!(q::Expr, ls::LoopSet, op::Operation, opp::Operation, par
   if isvectorized(opp)
     oppt = opp
   elseif isidentityop(opp)
-    oppt = only(parents(opp))
+    oppt = parents(opp)[1]
     # @show oppt
     isvectorized(oppt) || return parent
   else
@@ -570,7 +570,7 @@ function lower_compute!(
         else# mask last u₁
             :ifelselast # ifelse only the last one
         end
-        if last(instrcall.args) == varsym
+        if last(instrcall.args) === varsym
             pushfirst!(instrcall.args, lv(ifelsefunc))
             # showexpr = true
             insert!(instrcall.args, 3, MASKSYMBOL)

src/codegen/lower_load.jl

@@ -74,7 +74,7 @@ function add_prefetches!(q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, pr
     if !isknown(prefetchloop_step)
         for i ∈ eachindex(gespinds.args)
             if i == prefetchind
-                gespinds.args[i] = mulexpr(getsym(prefetchloop_step), gespinds.args[i])
+                gespinds.args[i] = mulexpr(getsym(prefetchloop_step), (gespinds.args[i])::Union{Symbol,Expr})
             end
             # gespinds.args[i] = Expr(:call, lv(:data), gespinds.args[i])
         end
@@ -163,31 +163,6 @@ function lower_load_no_optranslation!(
     end
     nothing
 end
-# function lower_load_vectorized!(
-#     q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned} = nothing
-# )
-#     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = td
-#     loopdeps = loopdependencies(op)
-#     @assert isvectorized(op)
-#     opu₁ = isu₁unrolled(op)
-#     inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
-#     if opu₁
-#         umin = 0
-#         U = u₁
-#     else
-#         umin = -1
-#         U = 0
-#     end
-#     # Urange = unrolled ∈ loopdeps ? 0:U-1 : 0
-#     var = variable_name(op, suffix)
-#     for u ∈ umin:U-1
-#         td = UnrollArgs(td, u)
-#         pushvectorload!(q, op, var, td, U, vectorized, mask, opu₁, inds_calc_by_ptr_offset, reg_size(ls))
-#     end
-#     prefetchind = prefetchisagoodidea(ls, op, td)
-#     iszero(prefetchind) || add_prefetches!(q, ls, op, td, prefetchind, umin)
-#     nothing
-# end
 function indisvectorized(ls::LoopSet, ind::Symbol)
     for op ∈ operations(ls)
         ((op.variable === ind) && isvectorized(op)) && return true
@@ -217,7 +192,6 @@ function lower_load_for_optranslation!(
     # abs of steps are equal
     equal_steps = (step₁ == step₂) ⊻ (posindicator ≠ 0x03)
     # @show step₁, step₂, posindicator, equal_steps
-    # _td = UnrollArgs(u₁loop, u₂loop, vloop, total_unroll, u₂max, Core.ifelse(equal_steps, 0, u₂max - 1))
     _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)
     ptr = vptr(op)

src/codegen/lower_memory_common.jl

@@ -40,22 +40,36 @@ function addoffset!(ret::Expr, stride, ind) # 3 args
 end
 # dropping `calcbypointeroffset` has to be delayed until after multiplying the `indexstride` by `index.
 function addoffset!(ret::Expr, stride, ind, offset, calcbypointeroffset::Bool) # 5 -> 3 args
-    if calcbypointeroffset
-        addoffset!(ret, stride, offset)
-    elseif _iszero(offset)
-        addoffset!(ret, stride, ind)
-    else
-        addoffset!(ret, stride, addexpr(ind, offset))
-    end
+  if calcbypointeroffset
+    addoffset!(ret, stride, offset)
+  elseif _iszero(offset)
+    addoffset!(ret, stride, ind)
+  else
+    addoffset!(ret, stride, addexpr(ind, offset))
+  end
+end
+function _addoffset!(ret::Expr, vloopstride, indexstride::Union{Integer,MaybeKnown}, index, offset, calcbypointeroffset::Bool) # 6 -> 5 args
+  if _isone(indexstride)
+    addoffset!(ret, vloopstride, index, offset, calcbypointeroffset)
+  else
+    __addoffset!(ret, vloopstride, indexstride, index, offset, calcbypointeroffset)
+  end
 end
 function _addoffset!(ret::Expr, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 6 -> 5 args
-    if _isone(indexstride)
-        addoffset!(ret, vloopstride, index, offset, calcbypointeroffset)
-    elseif isknown(vloopstride) & isknown(indexstride)
-        addoffset!(ret, gethint(vloopstride)*gethint(indexstride), index, offset, calcbypointeroffset)
-    else
-        addoffset!(ret, mulexpr(vloopstride,indexstride), index, offset, calcbypointeroffset)
-    end
+  ___addoffset!(ret, vloopstride, indexstride, index, offset, calcbypointeroffset)
+end
+function __addoffset!(ret::Expr, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 6 -> 5 args
+  ___addoffset!(ret, vloopstride, indexstride, index, offset, calcbypointeroffset)
+end
+function __addoffset!(ret::Expr, vloopstride::Union{Integer,MaybeKnown}, indexstride::Union{Integer,MaybeKnown}, index, offset, calcbypointeroffset::Bool) # 6 -> 5 args
+  if isknown(vloopstride) & isknown(indexstride)
+    addoffset!(ret, gethint(vloopstride)*gethint(indexstride), index, offset, calcbypointeroffset)
+  else
+    ___addoffset!(ret, vloopstride, indexstride, index, offset, calcbypointeroffset)
+  end
+end
+function ___addoffset!(ret::Expr, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 6 -> 5 args
+  addoffset!(ret, mulexpr(vloopstride,indexstride), index, offset, calcbypointeroffset)
 end
 # multiply `index` by `indexstride`
 function addoffset!(ret::Expr, vloopstride, indexstride, index, offset, calcbypointeroffset::Bool) # 6 -> (5 or 6) args
@@ -400,19 +414,19 @@ 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
+# 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)

src/codegen/lower_store.jl

@@ -1,12 +1,18 @@
+function opisreduced(op::Operation)
+  for rdep ∈ reduceddependencies(op)
+    rdep ∉ loopdependencies(op) && return true
+  end
+  false
+end
 function storeinstr_preprend(op::Operation, vloopsym::Symbol)
     # defaultstoreop = :vstore!
     # defaultstoreop = :vnoaliasstore!
     isvectorized(op) && return Symbol("")
     vloopsym ∉ reduceddependencies(op) && return Symbol("")
     # vectorized is not a loopdep, but is a reduced dep
-    opp = first(parents(op))
+    opp::Operation = first(parents(op))
     # while vectorized ∉ loopdependencies(opp)
-    while ((!isvectorized(opp)) || (any(rdep -> rdep ∉ loopdependencies(opp), reduceddependencies(opp))))
+    while ((!isvectorized(opp)) || opisreduced(opp))
         oppold = opp
         for oppp ∈ parents(opp)
             if vloopsym ∈ reduceddependencies(oppp)
@@ -145,88 +151,92 @@ function lower_store_collection!(
 end
 gf(s::Symbol, n::Int) = Expr(:call, GlobalRef(Core,:getfield), s, n, false)
 function lower_store!(
-    q::Expr, ls::LoopSet, op::Operation, ua::UnrollArgs, mask::Bool,
-    reductfunc::Symbol = storeinstr_preprend(op, ua.vloop.itersymbol), inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
-)
-    @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vloop, u₂max, suffix = ua
-    omop = offsetloadcollection(ls)
-    batchid, opind = omop.batchedcollectionmap[identifier(op)]
-    if ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op))
-        (opind == 1) && lower_store_collection!(q, ls, op, ua, mask, inds_calc_by_ptr_offset)
-        return
-    end
-    falseexpr = Expr(:call, lv(:False));
-    aliasexpr = falseexpr;
-    # trueexpr = Expr(:call, lv(:True));
-    rs = staticexpr(reg_size(ls));
-    opp = first(parents(op))
-    if (((opp.instruction.instr === reductfunc) || (opp.instruction.instr === :identity)) && isone(length(parents(opp))))
-        oppp = only(parents(opp))
-        if isu₂unrolled(op) == isu₂unrolled(oppp)
-            opp = oppp
-        end
+  q::Expr, ls::LoopSet, op::Operation, ua::UnrollArgs, mask::Bool,
+  reductfunc::Symbol = storeinstr_preprend(op, ua.vloop.itersymbol), inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
+  )
+  @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vloop, u₂max, suffix = ua
+  omop = offsetloadcollection(ls)
+  batchid, opind = omop.batchedcollectionmap[identifier(op)]
+  if ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op))
+    (opind == 1) && lower_store_collection!(q, ls, op, ua, mask, inds_calc_by_ptr_offset)
+    return
+  end
+  falseexpr = Expr(:call, lv(:False));
+  aliasexpr = falseexpr;
+  # trueexpr = Expr(:call, lv(:True));
+  rs = staticexpr(reg_size(ls));
+  opp = first(parents(op))
+  if ((opp.instruction.instr === reductfunc) || (opp.instruction.instr === :identity))
+    parents_opp = parents(opp)
+    opppstate = Base.iterate(parents_opp)
+    if opppstate ≢ nothing
+      oppp, state = opppstate
+      if (Base.iterate(parents_opp, state) === nothing) && isu₂unrolled(op) == isu₂unrolled(oppp)
+        opp = oppp
+      end
     end
+  end
   # __u₂max = ls.unrollspecification.u₂
-    isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, __u₂max)
-    # @show isu₁, isu₂, u₁loopsym, u₂loopsym
-    # @show isu₁, isu₂, opp, u₁loopsym, u₂loopsym, vloopsym
-    u = isu₁ ? u₁ : 1
-    mvar = Symbol(variable_name(opp, ifelse(isu₂, suffix, -1)), '_', u)
-    if all(op.ref.loopedindex)
-        inds = unrolledindex(op, ua, mask, inds_calc_by_ptr_offset, ls)
-        storeexpr = if reductfunc === Symbol("")
-            Expr(:call, lv(:_vstore!), sptr(op), mvar, inds)
-        else
-            Expr(:call, lv(:_vstore!), lv(reductfunc), sptr(op), mvar, inds)
-        end
-        add_memory_mask!(storeexpr, op, ua, mask, ls)
-        push!(storeexpr.args, falseexpr, aliasexpr, falseexpr, rs)
-        push!(q.args, storeexpr)
+  isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, __u₂max)
+  # @show isu₁, isu₂, u₁loopsym, u₂loopsym
+  # @show isu₁, isu₂, opp, u₁loopsym, u₂loopsym, vloopsym
+  u = isu₁ ? u₁ : 1
+  mvar = Symbol(variable_name(opp, ifelse(isu₂, suffix, -1)), '_', u)
+  if all(op.ref.loopedindex)
+    inds = unrolledindex(op, ua, mask, inds_calc_by_ptr_offset, ls)
+    storeexpr = if reductfunc === Symbol("")
+      Expr(:call, lv(:_vstore!), sptr(op), mvar, inds)
     else
-      parents_op = parents(op)
-      data_u₁ = isu₁ & (u₁ > 1)
-      
-      indices_u₁ = data_u₁
-      if !data_u₁ & (length(parents_op) > 1)
-        indices_u₁ = first(isunrolled_sym(op, u₁loopsym, u₂loopsym, vloopsym, ls))
-      end
-      if indices_u₁
-        mvard = Symbol(mvar, "##data##")
-        # isu₁ &&
-        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)
-          # @show isu₁unrolled(opp), opp
-          storeexpr = if data_u₁
-            if reductfunc === Symbol("")
-              Expr(:call, lv(:_vstore!), sptrsym, gf(mvard,u), inds)
-            else
-              Expr(:call, lv(:_vstore!), lv(reductfunc), sptrsym, mvaru, inds)
-            end
-          elseif reductfunc === Symbol("")
-            Expr(:call, lv(:_vstore!), sptrsym, mvar, inds)
+      Expr(:call, lv(:_vstore!), lv(reductfunc), sptr(op), mvar, inds)
+    end
+    add_memory_mask!(storeexpr, op, ua, mask, ls)
+    push!(storeexpr.args, falseexpr, aliasexpr, falseexpr, rs)
+    push!(q.args, storeexpr)
+  else
+    parents_op = parents(op)
+    data_u₁ = isu₁ & (u₁ > 1)
+    
+    indices_u₁ = data_u₁
+    if !data_u₁ & (length(parents_op) > 1)
+      indices_u₁ = first(isunrolled_sym(op, u₁loopsym, u₂loopsym, vloopsym, ls))
+    end
+    if indices_u₁
+      mvard = Symbol(mvar, "##data##")
+      # isu₁ &&
+      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)
+        # @show isu₁unrolled(opp), opp
+        storeexpr = if data_u₁
+          if reductfunc === Symbol("")
+            Expr(:call, lv(:_vstore!), sptrsym, gf(mvard,u), inds)
           else
-            Expr(:call, lv(:_vstore!), lv(reductfunc), sptrsym, mvar, inds)
+            Expr(:call, lv(:_vstore!), lv(reductfunc), sptrsym, mvaru, inds)
           end
-          domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
-          add_memory_mask!(storeexpr, op, ua, domask, ls)# & ((u == u₁) | isvectorized(op)))
-          push!(storeexpr.args, falseexpr, aliasexpr, falseexpr, rs)
-          push!(q.args, storeexpr)
-        end
-      else
-        inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, 0, ls)
-        storeexpr = if reductfunc === Symbol("")
-          Expr(:call, lv(:_vstore!), sptr(op), mvar, inds)
+        elseif reductfunc === Symbol("")
+          Expr(:call, lv(:_vstore!), sptrsym, mvar, inds)
         else
-          Expr(:call, lv(:_vstore!), lv(reductfunc), sptr(op), mvar, inds)
+          Expr(:call, lv(:_vstore!), lv(reductfunc), sptrsym, mvar, inds)
         end
-        add_memory_mask!(storeexpr, op, ua, mask, ls)
+        domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
+        add_memory_mask!(storeexpr, op, ua, domask, ls)# & ((u == u₁) | isvectorized(op)))
         push!(storeexpr.args, falseexpr, aliasexpr, falseexpr, rs)
         push!(q.args, storeexpr)
       end
+    else
+      inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, true, 0, ls)
+      storeexpr = if reductfunc === Symbol("")
+        Expr(:call, lv(:_vstore!), sptr(op), mvar, inds)
+      else
+        Expr(:call, lv(:_vstore!), lv(reductfunc), sptr(op), mvar, inds)
+      end
+      add_memory_mask!(storeexpr, op, ua, mask, ls)
+      push!(storeexpr.args, falseexpr, aliasexpr, falseexpr, rs)
+      push!(q.args, storeexpr)
     end
-    nothing
+  end
+  nothing
 end
 
 function lower_tiled_store!(

src/codegen/lowering.jl

@@ -677,7 +677,7 @@ function determine_eltype(ls::LoopSet, ortypdefined::Bool)::Union{Symbol,Expr}
     if narrays == 1
       return Expr(:call, lv(:eltype), first(ls.includedactualarrays))
     else
-      oreducop = ls.operations[only(ls.outer_reductions)]
+      oreducop = ls.operations[ls.outer_reductions[1]]
       if ortypdefined
         return typeof_expr(oreducop)
       else

src/codegen/operation_evaluation_order.jl

@@ -25,24 +25,28 @@
 
 function isnopidentity(ls::LoopSet, op::Operation, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, u₂max::Int)
   parents_op = parents(op)
-    if iscompute(op) && instruction(op).instr === :identity && isone(length(parents_op)) && name(first(parents_op)) === name(op)
-        # loopistiled = u₂max ≠ -1
-        # parents_u₁syms, parents_u₂syms = parent_unroll_status(op, u₁loop, u₂loop, u₂max)
-        # if (u₁unrolledsym == first(parents_u₁syms)) && (isu₂unrolled(op) == parents_u₂syms[1])
-        opp = only(parents_op)
-        # @show op opp isu₁unrolled(op), isu₁unrolled(opp), isu₂unrolled(op), isu₂unrolled(opp)
-        if (isu₁unrolled(op) == isu₁unrolled(opp)) & (isu₂unrolled(op) == isu₂unrolled(opp))
-            true
-        else
-            # if isvectorized(opp) & (!isvectorized(op))
-            #     op.instruction = reduction_to_scalar(instruction(opp))
-            #     op.mangledvariable = gensym(op.mangledvariable)
-            # end
-            false
-        end
+  if iscompute(op) && instruction(op).instr === :identity
+    # loopistiled = u₂max ≠ -1
+    # parents_u₁syms, parents_u₂syms = parent_unroll_status(op, u₁loop, u₂loop, u₂max)
+    # if (u₁unrolledsym == first(parents_u₁syms)) && (isu₂unrolled(op) == parents_u₂syms[1])
+    oppstate = Base.iterate(parents_op)
+    oppstate === nothing && return false
+    opp, state = oppstate
+    Base.iterate(parents_op, state) === nothing || return false
+    name(opp) === name(op) || return false
+    # @show op opp isu₁unrolled(op), isu₁unrolled(opp), isu₂unrolled(op), isu₂unrolled(opp)
+    if (isu₁unrolled(op) == isu₁unrolled(opp)) & (isu₂unrolled(op) == isu₂unrolled(opp))
+      true
     else
-        false
+      # if isvectorized(opp) & (!isvectorized(op))
+      #     op.instruction = reduction_to_scalar(instruction(opp))
+      #     op.mangledvariable = gensym(op.mangledvariable)
+      # end
+      false
     end
+  else
+    false
+  end
 end
 
 function set_upstream_family!(adal::Vector{T}, op::Operation, val::T, ld::Vector{Symbol}, id::Int) where {T}