Add support for storing loop induction variables directly, fixes #392.

Author: chriselrod

Project.toml

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

src/codegen/loopstartstopmanager.jl

@@ -92,7 +92,6 @@ function indices_calculated_by_pointer_offsets(ls::LoopSet, ar::ArrayReferenceMe
   ls.isbroadcast && return fill(false, length(indices))
   looporder = names(ls)
   offset = isdiscontiguous(ar)
-  gespinds = Expr(:tuple)
   out = Vector{Bool}(undef, length(indices))
   strds = getstrides(ar)
   li = ar.loopedindex
@@ -310,7 +309,7 @@ function isloopvalue(
     if (isrooted ≢ nothing)
       isrooted[i] || continue
     end
-    iscompute(op) || continue
+    (iscompute(op) | isstore(op)) || continue
     for opp ∈ parents(op)# this is to confirm `ind` still has children
       # (isloopvalue(opp) && instruction(opp).instr === ind) && return true
       if (isloopvalue(opp) && instruction(opp).instr === ind)
@@ -331,10 +330,8 @@ function cse_constant_offsets!(
   # 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
   li = ar.loopedindex
   indices = getindices(ar)
-  strides = getstrides(ar)
   offset = first(indices) === DISCONTIGUOUS
   # gespindoffsets = fill(Symbol(""), length(li))
   gespindsummary = Vector{Symbol}(undef, length(li))
@@ -814,7 +811,6 @@ function use_loop_induct_var!(
   q::Expr,
   ar::ArrayReferenceMeta,
   allarrayrefs::Vector{ArrayReferenceMeta},
-  allarrayrefsind::Int,
   includeinlet::Bool,
   # array_refs_with_same_name::Vector{Int}, arrayref_to_name_op_collection::Vector{Vector{Tuple{Int,Int,Int}}}
 )::Vector{Int}
@@ -868,7 +864,7 @@ function use_loop_induct_var!(
       Wisz || pushgespind!(gespinds, ls, Symbol(""), 0, 1, ind, isli, true, false)
     else
       uliv[i] = findfirst(Base.Fix2(===, ind), looporder)::Int
-      loop = getloop(ls, ind)
+      # 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, 1, ind, isli, false, false)
@@ -904,7 +900,6 @@ end
 # Plan here is that we increment every unique array
 function add_loop_start_stop_manager!(ls::LoopSet)
   q = Expr(:block)
-  us = ls.unrollspecification
   # Presence of an explicit use of a loopinducation var means we should use that, so we look for one
   # TODO: replace first with only once you add Compat as a dep or drop support for older Julia versions
   loopinductvars = Symbol[]
@@ -917,7 +912,7 @@ function add_loop_start_stop_manager!(ls::LoopSet)
   use_livs = Vector{Vector{Int}}(undef, length(arrayrefs))
   # for i ∈ eachindex(name_to_array_map)
   for i ∈ eachindex(arrayrefs)
-    use_livs[i] = use_loop_induct_var!(ls, q, arrayrefs[i], arrayrefs, i, includeinlet[i])
+    use_livs[i] = use_loop_induct_var!(ls, q, arrayrefs[i], arrayrefs, includeinlet[i])
     #name_to_array_map[first(first(unique_to_name_and_op_map[i]))], unique_to_name_and_op_map)
   end
   # loops, sorted from outer-most to inner-most
@@ -934,11 +929,9 @@ function add_loop_start_stop_manager!(ls::LoopSet)
   reached_indices = zeros(Int, length(arrayrefs))
   for (i, loopsym) ∈ enumerate(looporder) # iterates from outer to inner
     loopstartᵢ = ArrayReferenceMeta[]
-    arⱼ = 0
-    minrem = typemax(Int)
     ric = Tuple{Int,Int}[]
     for j ∈ eachindex(use_livs) # j is array ref number
-      for (l, k) ∈ enumerate(use_livs[j])# l is index number, k is loop number
+      for k ∈ use_livs[j]# l is index number, k is loop number
         if k == i
           push!(loopstartᵢ, arrayrefs[j])
           push!(ric, ((reached_indices[j] += 1), length(loopstartᵢ)))
@@ -1006,7 +999,6 @@ function pointermax_index(
   # @unpack u₁loopnum, u₂loopnum, vloopnum, u₁, u₂ = us
   loopsym = names(ls)[n]
   index = Expr(:tuple)
-  found_loop_sym = false
   ind = 0
   for (j, i) ∈ enumerate(getindicesonly(ar))
     if i === loopsym
@@ -1227,7 +1219,6 @@ function startloop(ls::LoopSet, us::UnrollSpecification, n::Int, staticinit::Boo
   termind = lssm.terminators[n]
   ptrdefs = lssm.incrementedptrs[n]
   loopstart = Expr(:block)
-  firstloop = n == num_loops(ls)
   for ar ∈ ptrdefs
     ptr_offset = vptr_offset(ar)
     push!(loopstart.args, Expr(:(=), ptr_offset, ptr_offset))
@@ -1288,7 +1279,7 @@ function incrementloopcounter!(
   if iszero(termind) # increment liv
     push!(q.args, incrementloopcounter(us, n, loopsym, UF, loop))
   end
-  for (j, ar) ∈ enumerate(ptrdefs)
+  for ar ∈ ptrdefs
     offsetinds = indices_calculated_by_pointer_offsets(ls, ar)
     push!(q.args, offset_ptr(ar, us, loopsym, n, UF, offsetinds, loop))
   end

src/codegen/lower_store.jl

@@ -94,8 +94,8 @@ function lower_store_collection!(
   inds_calc_by_ptr_offset::Vector{Bool},
 )
   omop = offsetloadcollection(ls)
-  batchid, bopind = omop.batchedcollectionmap[identifier(op)]
-  collectionid, copind = omop.opidcollectionmap[identifier(op)]
+  batchid, _ = omop.batchedcollectionmap[identifier(op)]
+  collectionid, __ = omop.opidcollectionmap[identifier(op)]
   opidmap = offsetloadcollection(ls).opids[collectionid]
   idsformap = omop.batchedcollections[batchid]
 
@@ -105,15 +105,18 @@ function lower_store_collection!(
   nouter = length(idsformap)
 
   t = Expr(:tuple)
-  for (i, (opid, _)) ∈ enumerate(idsformap)
+  for (opid, _) ∈ idsformap
     opp = first(parents(ops[opidmap[opid]]))
-
+    
     isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, __u₂max)
     u = Core.ifelse(isu₁, u₁, 1)
-    mvar = Symbol(variable_name(opp, ifelse(isu₂, suffix, -1)), '_', u)
-    # mvar = Symbol(variable_name(_op, suffix), '_', u)
-    # mvar = Symbol(variable_name(_op, ifelse(isu₂, suffix, -1)), '_', u)
-    push!(t.args, mvar)
+    if isloopvalue(opp)
+      loopval = first(loopdependencies(opp))
+      add_loopvalue!(t, loopval, ua, u₁, getloop(ls, loopval))
+    else
+      mvar = Symbol(variable_name(opp, ifelse(isu₂, suffix, -1)), '_', u)
+      push!(t.args, mvar)
+    end
   end
   offset_dummy_loop = Loop(
     first(getindices(op)),
@@ -149,7 +152,6 @@ function lower_store_collection!(
     false
   end
   uinds = Expr(:call, unrollcurl₂, inds)
-  vp = vptr(op)
   sptrsym = sptr!(q, op)
   storeexpr = Expr(:call, lv(:_vstore!), sptrsym, Expr(:call, lv(:VecUnroll), t), uinds)
   # not using `add_memory_mask!(storeexpr, op, ua, mask, ls)` because we checked `isconditionalmemop` earlier in `lower_load_collection!`
@@ -235,6 +237,12 @@ function lower_store!(
   isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, __u₂max)
   u = isu₁ ? u₁ : 1
   mvar = Symbol(variable_name(opp, ifelse(isu₂, suffix, -1)), '_', u)
+  if isloopvalue(opp)
+    def = Expr(:(=), mvar)
+    loopval = first(loopdependencies(opp))
+    add_loopvalue!(def, loopval, ua, u₁, getloop(ls, loopval))
+    push!(q.args, def)
+  end
   if all(op.ref.loopedindex)
     inds = unrolledindex(op, ua, mask, inds_calc_by_ptr_offset, ls)
     storeexpr = if reductfunc === Symbol("")
@@ -312,7 +320,6 @@ end
 function donot_tile_store(
   ls::LoopSet,
   op::Operation,
-  vloop::Loop,
   reductfunc::Symbol,
   u₂::Int,
 )
@@ -323,7 +330,7 @@ function donot_tile_store(
   ) && return true
   rejectcurly(op) && return true
   omop = offsetloadcollection(ls)
-  batchid, opind = omop.batchedcollectionmap[identifier(op)]
+  batchid, _ = omop.batchedcollectionmap[identifier(op)]
   return ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op))
 end
 
@@ -344,7 +351,7 @@ function lower_tiled_store!(
   reductfunc = storeinstr_preprend(op, vloopsym)
   inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
 
-  if donot_tile_store(ls, op, vloop, reductfunc, u₂)
+  if donot_tile_store(ls, op, reductfunc, u₂)
     # If we have a reductfunc, we're using a reducing store instead of a contiuguous or shuffle store anyway
     # so no benefit to being able to handle that case here, vs just calling the default `lower_store!` method
     @unpack u₁, u₂max = ua
@@ -368,6 +375,8 @@ function lower_tiled_store!(
   u = Core.ifelse(isu₁, u₁, 1)
   tup = Expr(:tuple)
   for t ∈ 0:u₂-1
+    # tiled stores cannot be loop values, as they're necessarilly
+    # functions of at least two loops, meaning we do not need to handle them here.
     push!(tup.args, Symbol(variable_name(opp, ifelse(isu₂, t, -1)), '_', u))
   end
   vut = Expr(:call, lv(:VecUnroll), tup) # `VecUnroll` of `VecUnroll`s

src/condense_loopset.jl

@@ -781,7 +781,6 @@ function generate_call_types(
   arrayref_descriptions = Expr(:tuple)
   duplicate_ref = fill(false, length(ls.refs_aliasing_syms))
   for (j, ref) ∈ enumerate(ls.refs_aliasing_syms)
-    vpref = vptr(ref)
     # duplicate_ref[j] ≠ 0 && continue
     duplicate_ref[j] && continue
     push!(arrayref_descriptions.args, ArrayRefStruct(ls, ref, arraysymbolinds, ids))

src/reconstruct_loopset.jl

@@ -980,7 +980,7 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
     ),
   )
   post = hoist_constant_memory_accesses!(ls)
-  # return @show avx_body(ls, var"#UNROLL#")
+  # q = @show(avx_body(ls, var"#UNROLL#")); post === ls.preamble ? q : Expr(:block, q, post)
   q = if (last(var"#UNROLL#") > 1) && length(var"#LPSYM#") == length(ls.loops)
     inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, l1, l2, l3, nt = var"#UNROLL#"
     # wrap in `var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#"` in `Expr` to homogenize types

test/copy.jl

@@ -183,7 +183,13 @@ using LoopVectorization, OffsetArrays, Test
       dest[i] = src[p, 3*p]
     end
   end
-
+  function collect_turbo(N, ::Type{T}) where {T}
+    l = Vector{T}(undef, N)
+    @turbo for i ∈ 1:length(l)
+      l[i] = i
+    end
+    return l
+  end
   for T ∈ (Float32, Float64, Int32, Int64)
     @show T, @__LINE__
     R = T <: Integer ? (-T(100):T(100)) : T
@@ -320,5 +326,7 @@ using LoopVectorization, OffsetArrays, Test
     end
     @. arr2[rng-ifirst] += 1
     @test arr1 == arr2
+
+    @test collect_turbo(77) == T.(1:77)
   end
 end