allow unroll curly in some more instances

Author: chriselrod

Project.toml

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

src/codegen/lower_load.jl

@@ -146,7 +146,7 @@ function lower_load_no_optranslation!(
   u = ifelse(opu₁, u₁, 1)
   mvar = Symbol(variable_name(op, Core.ifelse(opu₂, suffix,-1)), '_', u)
   falseexpr = Expr(:call, lv(:False)); rs = staticexpr(reg_size(ls))
-  if (all(op.ref.loopedindex) && !rejectcurly(op)) && vectorization_profitable(op)
+  if (!rejectcurly(op)) && vectorization_profitable(op)
     inds = unrolledindex(op, td, mask, inds_calc_by_ptr_offset, ls)
     loadexpr = Expr(:call, lv(:_vload), sptr(op), inds)
     add_memory_mask!(loadexpr, op, td, mask, ls, 0)
@@ -222,7 +222,7 @@ function lower_load_for_optranslation!(
     ip = GlobalRef(VectorizationBase, :increment_ptr)
     vpo = vptr_offset(gptr)
     push!(q.args, Expr(:(=), vpo, Expr(:call, ip, ptr, vptr_offset(ptr), gespinds)))
-    push!(q.args, Expr(:(=), gptr, ptr))#Expr(:call, GlobalRef(VectorizationBase, :reconstruct_ptr), 
+    push!(q.args, Expr(:(=), gptr, ptr))#Expr(:call, GlobalRef(VectorizationBase, :reconstruct_ptr),
     fill!(inds_by_ptroff, true)
     @unpack ref, loopedindex = mref
     indices = copy(getindices(ref))
@@ -367,9 +367,27 @@ function rejectcurly(ls::LoopSet, op::Operation, u₁loopsym::Symbol, vloopsym::
       end
     else
       opp = findop(parents(op), ind)
+      (isu₁unrolled(opp) || isu₂unrolled(opp)) && return true
+      length(parents(opp)) == 2 || return true
+      if instruction(opp).instr === :(+) || instruction(opp).instr === :add_fast
+        isadd = true
+      elseif instruction(opp).instr === :(-) || instruction(opp).instr === :sub_fast
+        isadd = false
+      else
+        return true
+      end
+      opp1 = parents(opp)[1]
+      opp2 = parents(opp)[2]
       if isvectorized(opp)
         AV && return true
         AV = true
+        if isvectorized(opp1)
+          isvectorized(opp2) && return true
+          isloopvalue(opp1) || return true
+        else# opp2 vectorized
+          isadd || return true
+          isloopvalue(opp2) || return true
+        end
       end
       if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp))
         AU && return true
@@ -393,7 +411,7 @@ function rejectinterleave(ls::LoopSet, op::Operation, vloop::Loop, idsformap::Su
       end
     end
   end
-  vloopsym = vloop.itersymbol; 
+  vloopsym = vloop.itersymbol;
   (first(getindices(op)) === vloopsym) && (length(idsformap) ≠ first(getstrides(op)) * gethint(strd))
 end
 # function lower_load_collection_manual_u₁unroll!(
@@ -501,12 +519,11 @@ function lower_load_collection!(
         u = Core.ifelse(opu₁, u₁, 1)
         for (i,(opid,o)) ∈ enumerate(idsformap)
             extractedv = Expr(:call, gf, collectionname, i, false)
-            
+
             _op = ops[opidmap[opid]]
             mvar = Symbol(variable_name(_op, Core.ifelse(opu₂, suffix, -1)), '_', u)
             push!(q.args, Expr(:(=), mvar, extractedv))
         end
         # unpack_collection!(q, ls, opidmap, idsformap, ua, loadexpr, collectionname, op, true)
     end
 end
-

src/modeling/determinestrategy.jl

@@ -43,7 +43,10 @@ function unitstride(ls::LoopSet, op::Operation, s::Symbol)
     end
     true
 end
-
+function cannot_shuffle(op::Operation, u₁::Symbol, u₂::Symbol, contigind::Symbol, indices) # assumes isvectorized and !unitstride
+  !((!rejectcurly(op) && (((contigind === CONSTANTZEROINDEX) && ((length(indices) > 1) && (indices[2] === u₁) || (indices[2] === u₂))) ||
+    ((u₁ === contigind) | (u₂ === contigind)))))
+end
 function cost(ls::LoopSet, op::Operation, (u₁,u₂)::Tuple{Symbol,Symbol}, vloopsym::Symbol, Wshift::Int, size_T::Int = op.elementbytes)
   isconstant(op) && return 0.0, 0, 1.0#Float64(length(loopdependencies(op)) > 0)
   isloopvalue(op) && return 0.0, 0, 0.0
@@ -67,22 +70,36 @@ function cost(ls::LoopSet, op::Operation, (u₁,u₂)::Tuple{Symbol,Symbol}, vlo
         indices = getindices(op)
         contigind = first(indices)
         shifter = max(2,Wshift)
-        if rejectinterleave(op)
-          offset = 0.0 # gather/scatter, alignment doesn't matter
-        else
-          shifter -= 1
-          offset = 0.5reg_size(ls) / cache_lnsze(ls)
-          if shifter > 1 &&
-            (!rejectcurly(op) && (((contigind === CONSTANTZEROINDEX) && ((length(indices) > 1) && (indices[2] === u₁) || (indices[2] === u₂))) ||
-            ((u₁ === contigind) | (u₂ === contigind))))
-
-            shifter -= 1
-            offset = 0.5reg_size(ls) / cache_lnsze(ls)
+        # rejectinterleave false means omop
+        # cannot shuffle false means reject curly
+        # either false means shuffle
+        dont_shuffle = rejectinterleave(op) && cannot_shuffle(op, u₁, u₂, contigind, indices)
+        if dont_shuffle
+          # offset = 0.0 # gather/scatter, alignment doesn't matter
+          r = 1 << shifter
+          srt = srt*r# + offset
+          sl *= r
+        else#if rejectinterleave(op) # means omop
+          if isload(op) & (length(loopdependencies(op)) > 1)# vmov(a/u)pd
+            srt += 0.5reg_size(ls) / cache_lnsze(ls)
           end
+          # srt += 0.3shifter # shifter == number of shuffles
+          # sl += 0.3shifter
+          srt += shifter # shifter == number of shuffles
+          sl += shifter
+          # shifter -= 1
+          # offset = 0.5reg_size(ls) / cache_lnsze(ls)
+          # r = 1 << shifter
+          # srt = srt*r + offset
+          # sl *= r
+        #   if shifter > 1 && (!(cannot_shuffle(op, u₁, u₂, contigind, indices)))
+        #     shifter -= 1
+        #     offset = 0.5reg_size(ls) / cache_lnsze(ls)
+        #   end
+        # else
         end
-        r = 1 << shifter
-        srt = srt*r + offset
-        sl *= r
+        # @show srt, sl
+        # @show shifter, offset, dont_shuffle
       elseif isload(op) & (length(loopdependencies(op)) > 1)# vmov(a/u)pd
         # penalize vectorized loads with more than 1 loopdep
         # heuristic; more than 1 loopdep means that many loads will not be aligned
@@ -94,7 +111,7 @@ function cost(ls::LoopSet, op::Operation, (u₁,u₂)::Tuple{Symbol,Symbol}, vlo
         srt += 0.5reg_size(ls) / cache_lnsze(ls)
         # srt += 0.25reg_size(ls) / cache_lnsze(ls)
       end
-    elseif isstore(op) # broadcast or reductionstore; if store we want to penalize reduction
+    elseif isstore(op)# && isvectorized(first(parents(op))) # broadcast or reductionstore; if store we want to penalize reduction
       srt *= 3
       sl *= 3
     end
@@ -828,6 +845,7 @@ function load_elimination_cost_factor!(
     # cost_vec[1] -= 0.5625 * iters
     # cost_vec[1] -= 0.5625 * iters / 2
     # @show rto, 0.8rt, op
+    # reg_pressure[1] += 0.25rp
     reg_pressure[1] += 0.25rp
     cost_vec[2] += rt
     reg_pressure[2] += rp

src/modeling/graphs.jl

@@ -570,21 +570,24 @@ function fill_children!(ls::LoopSet)
     end
   end
 end
+function rejectinterleave!(ls::LoopSet, op::Operation, u₁loop::Symbol, u₂loop::Symbol, vloopsym::Symbol, vloop::Loop)
+  setunrolled!(ls, op, u₁loop, u₂loop, vloopsym)
+  if accesses_memory(op)
+    rc = rejectcurly(ls, op, u₁loop, vloopsym)
+    op.rejectcurly = rc
+    if rc
+      op.rejectinterleave = true
+    else
+      omop = ls.omop
+      batchid, opind = omop.batchedcollectionmap[identifier(op)]
+      op.rejectinterleave = ((batchid == 0) || (!isvectorized(op))) || rejectinterleave(ls, op, vloop, omop.batchedcollections[batchid])
+    end
+  end
+end
 function cacheunrolled!(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, vloopsym::Symbol)
   vloop = getloop(ls, vloopsym)
   for op ∈ operations(ls)
-    setunrolled!(ls, op, u₁loop, u₂loop, vloopsym)
-    if accesses_memory(op)
-      rc = rejectcurly(ls, op, u₁loop, vloopsym)
-      op.rejectcurly = rc
-      if rc
-        op.rejectinterleave = true
-      else
-        omop = ls.omop
-        batchid, opind = omop.batchedcollectionmap[identifier(op)]
-        op.rejectinterleave = ((batchid == 0) || (!isvectorized(op))) || rejectinterleave(ls, op, vloop, omop.batchedcollections[batchid])
-      end
-    end
+    rejectinterleave!(ls, op, u₁loop, u₂loop, vloopsym, vloop)
   end
 end
 function setunrolled!(ls::LoopSet, op::Operation, u₁loopsym::Symbol, u₂loopsym::Symbol, vectorized::Symbol)
@@ -1411,7 +1414,6 @@ function fill_offset_memop_collection!(ls::LoopSet)
   end
   for (collectionid,opidc) ∈ enumerate(opids)
     length(opidc) > 1 || continue
-
     # we check if we can turn the offsets into an unroll
     # we have up to `length(opidc)` loads to do, so we allocate that many "base" vectors
     # then we iterate through them, adding them to collections as appropriate

src/reconstruct_loopset.jl

@@ -711,7 +711,7 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
 - `vargs...` holds the encoded pointers of all the arrays (see `VectorizationBase`'s various pointer types).
 """
 @generated function _turbo_!(
-    ::Val{var"#UNROLL#"}, ::Val{var"#OPS#"}, ::Val{var"#ARF#"}, ::Val{var"#AM#"}, ::Val{var"#LPSYM#"}, ::Val{Tuple{var"#LB#",var"#V#"}}, var"#flattened#var#arguments#"::Vararg{Any,var"#num#vargs#"}
+  ::Val{var"#UNROLL#"}, ::Val{var"#OPS#"}, ::Val{var"#ARF#"}, ::Val{var"#AM#"}, ::Val{var"#LPSYM#"}, ::Val{Tuple{var"#LB#",var"#V#"}}, var"#flattened#var#arguments#"::Vararg{Any,var"#num#vargs#"}
 ) where {var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#", var"#V#", var"#num#vargs#"}
   # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
   ls = _turbo_loopset(var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#".parameters, var"#V#".parameters, var"#UNROLL#")