Add constant offset indices based dependency check

Author: chriselrod

Project.toml

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

src/codegen/lower_load.jl

@@ -331,18 +331,20 @@ end
 function _lower_load!(
     q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool} = indices_calculated_by_pointer_offsets(ls, op.ref)
 )
-    if rejectinterleave(op)
-        lower_load_no_optranslation!(q, ls, op, td, mask, inds_calc_by_ptr_offset)
-    else
-        omop = offsetloadcollection(ls)
-        batchid, opind = omop.batchedcollectionmap[identifier(op)]
-        if opind == 1
-            collectionid, copind = omop.opidcollectionmap[identifier(op)]
-            opidmap = offsetloadcollection(ls).opids[collectionid]
-            idsformap = omop.batchedcollections[batchid]
-            lower_load_collection!(q, ls, opidmap, idsformap, td, mask, inds_calc_by_ptr_offset)
-        end
+  if rejectinterleave(op)
+    lower_load_no_optranslation!(q, ls, op, td, mask, inds_calc_by_ptr_offset)
+  else
+    omop = offsetloadcollection(ls)
+    @unpack opids, opidcollectionmap, batchedcollections, batchedcollectionmap = omop
+    batchid, opind = batchedcollectionmap[identifier(op)]
+    if opind == 1
+      collectionid, copind = opidcollectionmap[identifier(op)]
+      opidmap = opids[collectionid]
+      idsformap = batchedcollections[batchid]
+      lower_load_collection!(q, ls, opidmap, idsformap, td, mask, inds_calc_by_ptr_offset)
     end
+  end
+  return nothing
 end
 function additive_vectorized_loopinductvar_only(op::Operation)
     isvectorized(op) || return true
@@ -388,18 +390,21 @@ function rejectcurly(ls::LoopSet, op::Operation, u₁loopsym::Symbol, vloopsym::
     false
 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)
-                additive_vectorized_loopinductvar_only(indop) || return true # so that it is `MM`
-            end
-        end
+  strd = step(vloop)
+  isknown(strd) || return true
+  # TODO: reject if there is a vectorized !loopedindex index
+  indices = getindicesonly(op); li = op.ref.loopedindex
+  for i ∈ eachindex(li)
+    li[i] && continue
+    ind = indices[i]
+    for indop ∈ operations(ls)
+      if (name(indop) === ind) && isvectorized(indop)
+        additive_vectorized_loopinductvar_only(indop) || return true # so that it is `MM`
+      end
     end
-    (first(getindices(op)) === vloopsym) && (length(idsformap) ≠ first(getstrides(op)) * gethint(strd))
+  end
+  vloopsym = vloop.itersymbol; 
+  (first(getindices(op)) === vloopsym) && (length(idsformap) ≠ first(getstrides(op)) * gethint(strd))
 end
 # function lower_load_collection_manual_u₁unroll!(
 #     q::Expr, ls::LoopSet, opidmap::Vector{Int},

src/codegen/lower_memory_common.jl

@@ -185,80 +185,83 @@ end
 
 # interleave: `0` means `false`, positive means literal, negative means multiplier
 function unrolled_curly(op::Operation, u₁::Int, u₁loop::Loop, vloop::Loop, mask::Bool, interleave::Int=0)
-    u₁loopsym = u₁loop.itersymbol
-    vloopsym = vloop.itersymbol
-    indices = getindicesonly(op)
-    vstep = step(vloop)
-    li = op.ref.loopedindex
-    # @assert all(loopedindex)
-    # @unpack u₁, u₁loopsym, vloopsym = td
-    AV = AU = -1
-    for (n,ind) ∈ enumerate(indices)
-        # @show AU, op, n, ind, vloopsym, u₁loopsym
-        if li[n]
-            if ind === vloopsym
-                @assert AV == -1 # FIXME: these asserts should be replaced with checks that prevent using `unrolled_curly` in these cases (also to be reflected in cost modeling, to avoid those)
-                AV = n
-            end
-            if ind === u₁loopsym
-                @assert AU == -1
-                AU = n
-            end
-        else
-            opp = findop(parents(op), ind)
-            # @show opp
-            if isvectorized(opp)
-                @assert AV == -1
-                AV = n
-            end
-            if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp) || (ind === u₁loopsym))
-                @assert AU == -1
-                AU = n
-            end
-        end
-    end
-    AU == -1 && throw(LoopError("Failed to find $(u₁loopsym) in args of $(repr(op))."))
-    vecnotunrolled = AU != AV
-    conditional_memory_op = isconditionalmemop(op)
-    if mask || conditional_memory_op
-        M = one(UInt)
-        # `isu₁unrolled(last(parents(op)))` === is condop unrolled?
-        # isu₁unrolled(last(parents(op)))
-        if vecnotunrolled || conditional_memory_op || (interleave > 0) # mask all
-            M = (M << u₁) - M
-        else # mask last
-            M <<= (u₁ - 1)
+  u₁loopsym = u₁loop.itersymbol
+  vloopsym = vloop.itersymbol
+  indices = getindicesonly(op)
+  vstep = step(vloop)
+  li = op.ref.loopedindex
+  # @assert all(loopedindex)
+  # @unpack u₁, u₁loopsym, vloopsym = td
+  AV = AU = -1
+  for (n,ind) ∈ enumerate(indices)
+    # @show AU, op, n, ind, vloopsym, u₁loopsym
+    if li[n]
+      if ind === vloopsym
+        @assert AV == -1 # FIXME: these asserts should be replaced with checks that prevent using `unrolled_curly` in these cases (also to be reflected in cost modeling, to avoid those)
+        AV = n
+      end
+      if ind === u₁loopsym
+        if AU ≠ -1
+          u₁loopsym === CONSTANTZEROINDEX && continue
+          throw(ArgumentError("Two of the same index $ind?"))
         end
+        AU = n
+      end
     else
-        M = zero(UInt)
+      opp = findop(parents(op), ind)
+      # @show opp
+      if isvectorized(opp)
+        @assert AV == -1
+        AV = n
+      end
+      if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp) || (ind === u₁loopsym))
+        @assert AU == -1
+        AU = n
+      end
+    end
+  end
+  AU == -1 && throw(LoopError("Failed to find $(u₁loopsym) in args of $(repr(op))."))
+  vecnotunrolled = AU != AV
+  conditional_memory_op = isconditionalmemop(op)
+  if mask || conditional_memory_op
+    M = one(UInt)
+    # `isu₁unrolled(last(parents(op)))` === is condop unrolled?
+    # isu₁unrolled(last(parents(op)))
+    if vecnotunrolled || conditional_memory_op || (interleave > 0) # mask all
+      M = (M << u₁) - M
+    else # mask last
+      M <<= (u₁ - 1)
     end
-    @assert isknown(step(u₁loop)) "Unrolled loops must have known steps to use `Unroll` type; this is a bug, shouldn't have reached here"
-    if AV > 0
-        @assert isknown(step(vloop)) "Vectorized loops must have known steps to use `Unroll` type; this is a bug, shouldn't have reached here."
-        X = convert(Int, getstrides(op)[AV])
-        X *= gethint(step(vloop))
-        intvecsym = :(Int($VECTORWIDTHSYMBOL))
-        if interleave > 0
-            Expr(:curly, lv(:Unroll), AU, interleave, u₁, AV, intvecsym, M, X)
-        elseif interleave < 0
-            unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, -interleave))))
-            Expr(:curly, lv(:Unroll), AU, unrollstepexpr, u₁, AV, intvecsym, M, X)
+  else
+    M = zero(UInt)
+  end
+  @assert isknown(step(u₁loop)) "Unrolled loops must have known steps to use `Unroll` type; this is a bug, shouldn't have reached here"
+  if AV > 0
+    @assert isknown(step(vloop)) "Vectorized loops must have known steps to use `Unroll` type; this is a bug, shouldn't have reached here."
+    X = convert(Int, getstrides(op)[AV])
+    X *= gethint(step(vloop))
+    intvecsym = :(Int($VECTORWIDTHSYMBOL))
+    if interleave > 0
+      Expr(:curly, lv(:Unroll), AU, interleave, u₁, AV, intvecsym, M, X)
+    elseif interleave < 0
+      unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, -interleave))))
+      Expr(:curly, lv(:Unroll), AU, unrollstepexpr, u₁, AV, intvecsym, M, X)
+    else
+      if vecnotunrolled
+        # 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(X)
+          Expr(:curly, lv(:Unroll), AU, intvecsym, u₁, AV, intvecsym, M, X)
         else
-            if vecnotunrolled
-                # 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(X)
-                    Expr(:curly, lv(:Unroll), AU, intvecsym, u₁, AV, intvecsym, M, X)
-                else
-                    unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, X))))
-                    Expr(:curly, lv(:Unroll), AU, unrollstepexpr, u₁, AV, intvecsym, M, X)
-                end
-            end
+          unrollstepexpr = :(Int($(mulexpr(VECTORWIDTHSYMBOL, X))))
+          Expr(:curly, lv(:Unroll), AU, unrollstepexpr, u₁, AV, intvecsym, M, X)
         end
-    else
-        Expr(:curly, lv(:Unroll), AU, gethint(step(u₁loop)), u₁, 0, 1, M, 1)
+      end
     end
+  else
+    Expr(:curly, lv(:Unroll), AU, gethint(step(u₁loop)), u₁, 0, 1, M, 1)
+  end
 end
 function unrolledindex(op::Operation, td::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool}, ls::LoopSet)
     @unpack u₁, u₁loopsym, u₁loop, vloop = td

src/codegen/lowering.jl

@@ -26,7 +26,9 @@ function lower!(
         end
     end
 end
-
+function isu₂invalidstorereorder(ls::LoopSet, us::UnrollSpecification)
+  us.u₂ == -1 ? false : ls.validreorder[ls.loopordermap[us.u₂loopnum]] ≠ 0x03
+end
 function lower_block(
     ls::LoopSet, us::UnrollSpecification, n::Int, mask::Bool, UF::Int
 )
@@ -40,6 +42,7 @@ function lower_block(
     u₁ = n == u₁loopnum ? UF : u₁
     dontmaskfirsttiles = mask && vloopnum == u₂loopnum
     blockq = Expr(:block)
+    cannot_reorder_u₂ = isu₂invalidstorereorder(ls, us)
     for prepost ∈ 1:2
         # !u₁ && !u₂
         lower!(blockq, ops[1,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, true)
@@ -56,22 +59,30 @@ function lower_block(
                 lower_tiled_store!(blockq, opsv1, opsv2, ls, unrollsyms, u₁, u₂, mask)
             else
                 for store ∈ (false,true)
+                    if cannot_reorder_u₂
+                        nstores = 0# break
+                        lowernonstore = lowerstore = true
+                    else
+                        lowerstore = store; lowernonstore = !store
+                    end
                     for t ∈ 0:u₂-1
                         # !u₁ &&  u₂
-                        lower!(blockq, opsv1, ls, unrollsyms, u₁, u₂, t, mask & !(dontmaskfirsttiles & (t < u₂ - 1)), !store, store)
+                        lower!(blockq, opsv1, ls, unrollsyms, u₁, u₂, t, mask & !(dontmaskfirsttiles & (t < u₂ - 1)), lowernonstore, lowerstore)
                         if iszero(t) && !store #  u₁ && !u₂
                             # for u ∈ 0:u₁-1
-                            lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, true)
+                            lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, lowernonstore, lowerstore)
                             # end
                         end
                         #  u₁ && u₂
                         # for u ∈ 0:u₁-1
-                        lower!(blockq, opsv2, ls, unrollsyms, u₁, u₂, t, mask & !(dontmaskfirsttiles & (t < u₂ - 1)), !store, store)
+                        lower!(blockq, opsv2, ls, unrollsyms, u₁, u₂, t, mask & !(dontmaskfirsttiles & (t < u₂ - 1)), lowernonstore, lowerstore)
                         # end
                     end
                     nstores == 0 && break
                 end
             end
+        elseif cannot_reorder_u₂
+            lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, true)
         else
             # for u ∈ 0:u₁-1     #  u₁ && !u₂
             lower!(blockq, ops[2,1,prepost,n], ls, unrollsyms, u₁, u₂, -1, mask, true, false)