hacky fix for CartesianIndices{0} stores that works for some cases

Author: chriselrod

src/LoopVectorization.jl

@@ -91,6 +91,7 @@ include("codegen/lowering.jl")
 include("codegen/split_loops.jl")
 include("codegen/lower_threads.jl")
 include("condense_loopset.jl")
+include("transforms.jl")
 include("reconstruct_loopset.jl")
 include("constructors.jl")
 include("user_api_conveniences.jl")

src/codegen/loopstartstopmanager.jl

@@ -37,35 +37,36 @@ function uniquearrayrefs_csesummary(ls::LoopSet)
 end
 
 function uniquearrayrefs(ls::LoopSet)
-    uniquerefs = ArrayReferenceMeta[]
-    includeinlet = Bool[]
-    # for arrayref ∈ ls.refs_aliasing_syms
-    for op ∈ operations(ls)
-        arrayref = op.ref
-        arrayref === NOTAREFERENCE && continue
-        notunique = false
-        isonlyname = true
-        for ref ∈ uniquerefs
-            notunique = sameref(arrayref, ref)
-            isonlyname &= vptr(arrayref) !== vptr(ref)
-            # if they're not the sameref, they may still have the same name
-            # if they have different names, they're definitely not sameref
-            notunique && break
-        end
-        if !notunique
-            push!(uniquerefs, arrayref)
-            push!(includeinlet, isonlyname)
-        end
+  uniquerefs = ArrayReferenceMeta[]
+  includeinlet = Bool[]
+  # for arrayref ∈ ls.refs_aliasing_syms
+  for op ∈ operations(ls)
+    op.instruction === DROPPEDCONSTANT && continue
+    arrayref = op.ref
+    arrayref === NOTAREFERENCE && continue
+    notunique = false
+    isonlyname = true
+    for ref ∈ uniquerefs
+      notunique = sameref(arrayref, ref)
+      isonlyname &= vptr(arrayref) !== vptr(ref)
+      # if they're not the sameref, they may still have the same name
+      # if they have different names, they're definitely not sameref
+      notunique && break
     end
-    uniquerefs, includeinlet
+    if !notunique
+      push!(uniquerefs, arrayref)
+      push!(includeinlet, isonlyname)
+    end
+  end
+  uniquerefs, includeinlet
 end
 
 otherindexunrolled(loopsym::Symbol, ind::Symbol, loopdeps::Vector{Symbol}) = (loopsym !== ind) && (loopsym ∈ loopdeps)
 function otherindexunrolled(ls::LoopSet, ind::Symbol, ref::ArrayReferenceMeta)
-    us = ls.unrollspecification
-    u₁sym = names(ls)[us.u₁loopnum]
-    u₂sym = us.u₂loopnum > 0 ? names(ls)[us.u₂loopnum] : Symbol("##undefined##")
-    otherindexunrolled(u₁sym, ind, loopdependencies(ref)) || otherindexunrolled(u₂sym, ind, loopdependencies(ref))
+  us = ls.unrollspecification
+  u₁sym = names(ls)[us.u₁loopnum]
+  u₂sym = us.u₂loopnum > 0 ? names(ls)[us.u₂loopnum] : Symbol("##undefined##")
+  otherindexunrolled(u₁sym, ind, loopdependencies(ref)) || otherindexunrolled(u₂sym, ind, loopdependencies(ref))
 end
 function multiple_with_name(n::Symbol, v::Vector{ArrayReferenceMeta})
   found = false
@@ -79,26 +80,26 @@ end
 # multiple_with_name(n::Symbol, v::Vector{ArrayReferenceMeta}) = sum(ref -> n === vptr(ref), v) > 1
 # TODO: DRY between indices_calculated_by_pointer_offsets and use_loop_induct_var
 function indices_calculated_by_pointer_offsets(ls::LoopSet, ar::ArrayReferenceMeta)
-    indices = getindices(ar)
-    ls.isbroadcast && return fill(false, length(indices))
-    looporder = names(ls)
-    offset = isdiscontiguous(ar)
-    gespinds = Expr(:tuple)
-    out = Vector{Bool}(undef, length(indices))
-    li = ar.loopedindex
-    for i ∈ eachindex(li)
-        ii = i + offset
-        ind = indices[ii]
-        if (!li[i]) || (ind === CONSTANTZEROINDEX) || multiple_with_name(vptr(ar), ls.lssm.uniquearrayrefs) ||
-            (iszero(ls.vector_width) && isstaticloop(getloop(ls, ind)))# ||
-            out[i] = false
-        elseif (isone(ii) && (first(looporder) === ind))
-            out[i] = otherindexunrolled(ls, ind, ar)
-        else
-            out[i] = true
-        end
+  indices = getindices(ar)
+  ls.isbroadcast && return fill(false, length(indices))
+  looporder = names(ls)
+  offset = isdiscontiguous(ar)
+  gespinds = Expr(:tuple)
+  out = Vector{Bool}(undef, length(indices))
+  li = ar.loopedindex
+  for i ∈ eachindex(li)
+    ii = i + offset
+    ind = indices[ii]
+    if (!li[i]) || (ind === CONSTANTZEROINDEX) || multiple_with_name(vptr(ar), ls.lssm.uniquearrayrefs) ||
+      (iszero(ls.vector_width) && isstaticloop(getloop(ls, ind)))# ||
+      out[i] = false
+    elseif (isone(ii) && (first(looporder) === ind))
+      out[i] = otherindexunrolled(ls, ind, ar)
+    else
+      out[i] = true
     end
-    out
+  end
+  out
 end
 
 # @generated function set_first_stride(sptr::StridedPointer{T,N,C,B,R}) where {T,N,C,B,R}
@@ -139,8 +140,8 @@ end
 
 # end
 function set_ref_loopedindex_and_ind!(ref::ArrayReferenceMeta, i::Int, ii::Int, li::Bool, ind::Symbol)
-    ref.loopedindex[i] = li
-    getindices(ref)[ii] = ind
+  ref.loopedindex[i] = li
+  getindices(ref)[ii] = ind
 end
 function set_all_to_constant_index!(
   ls::LoopSet, i::Int, ii::Int, indop::Operation, allarrayrefs::Vector{ArrayReferenceMeta},

src/condense_loopset.jl

@@ -741,6 +741,7 @@ make_crashy(q) = Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,
 
 @inline vecmemaybe(x::NativeTypes) = x
 @inline vecmemaybe(x::VectorizationBase._Vec) = Vec(x)
+@inline vecmemaybe(x::VectorizationBase.Vec) = x
 @inline vecmemaybe(x::Tuple) = VectorizationBase.VecUnroll(x)
 @inline vecmemaybe(x::Mask) = x
 
@@ -770,28 +771,31 @@ end
 #   setup_call_ret!(ls, call, preserve)
 # end
 setup_outerreduct_preserve_mangler(op::Operation) = Symbol(mangledvar(op), "##onevec##")
+
+function outer_reduction_to_scalar_reduceq!(q::Expr, op::Operation, var = name(op))
+  instr = instruction(op)
+  out = setup_outerreduct_preserve_mangler(op)
+  if instr.instr ≢ :ifelse
+    Expr(:call, reduction_scalar_combine(op), Expr(:call, lv(:vecmemaybe), out), var)
+  else
+    opinstr = ifelse_reduction(:IfElseReduced, op) do opv
+      opvname = name(opv)
+      oporig = gensym(opvname)
+      pushfirst!(q.args, Expr(:(=), oporig, opvname))
+      Expr(:call, lv(:vecmemaybe), setup_outerreduct_preserve_mangler(opv)), (oporig,)
+    end
+    Expr(:call, opinstr, Expr(:call, lv(:vecmemaybe), out), var)
+  end
+end
 function setup_outerreduct_preserve(ls::LoopSet, call::Expr, preserve::Vector{Symbol})
   iszero(length(ls.outer_reductions)) && return gc_preserve(call, preserve)
   retv = loopset_return_value(ls, Val(false))
   q = Expr(:block, gc_preserve(Expr(:(=), retv, call), preserve))
   for or ∈ ls.outer_reductions
     op = ls.operations[or]
-    var = name(op)
     # push!(call.args, Symbol("##TYPEOF##", var))
-    instr = instruction(op)
-    out = setup_outerreduct_preserve_mangler(op)
-    reducq = if instr.instr ≢ :ifelse
-      Expr(:call, reduction_scalar_combine(op), Expr(:call, lv(:vecmemaybe), out), var)
-    else
-      opinstr = ifelse_reduction(:IfElseReduced, op) do opv
-        opvname = name(opv)
-        oporig = gensym(opvname)
-        pushfirst!(q.args, Expr(:(=), oporig, opvname))
-        Expr(:call, lv(:vecmemaybe), setup_outerreduct_preserve_mangler(opv)), (oporig,)
-      end
-      Expr(:call, opinstr, Expr(:call, lv(:vecmemaybe), out), var)
-    end
-    push!(q.args, Expr(:(=), var, reducq))
+    reducq = outer_reduction_to_scalar_reduceq!(q, op)
+    push!(q.args, Expr(:(=), name(op), reducq))
   end
   q
 end

src/modeling/operations.jl

@@ -375,19 +375,18 @@ getstrides(ref) = arrayref(ref).strides
 
 isdiscontiguous(ref) = isdiscontiguous_inds(getindices(ref))
 function isdiscontiguous_inds(inds)
-    # (first(inds) === DISCONTIGUOUS) || (first(inds) === CONSTANTZEROINDEX)
-    first(inds) === DISCONTIGUOUS
+  length(inds) == 0 ? false : (@inbounds(inds[begin]) === DISCONTIGUOUS)
 end
 function makediscontiguous!(inds)
-    if iszero(length(inds)) || !isdiscontiguous_inds(inds)
-        pushfirst!(inds, DISCONTIGUOUS)
-    end
-    nothing
+  if iszero(length(inds)) || !isdiscontiguous_inds(inds)
+    pushfirst!(inds, DISCONTIGUOUS)
+  end
+  nothing
 end
 
 function getindicesonly(ref)
-    indices = getindices(ref)
-    @view(indices[isdiscontiguous(ref) + 1:end])
+  indices = getindices(ref)
+  @view(indices[isdiscontiguous(ref) + 1:end])
 end
 # function hasintersection(s1::Set{T}, s2::Set{T}) where {T}
     # for x ∈ s1

src/parse/add_constants.jl

@@ -99,9 +99,7 @@ function ensure_constant_lowered!(ls::LoopSet, mpref::ArrayReferenceMetaPosition
   end
   return nothing
 end
-function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementbytes::Int)
-  op = Operation(length(operations(ls)), varname(mpref), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS, mpref.mref)
-  add_vptr!(ls, op)
+function add_constant_vload!(ls::LoopSet, op::Operation, mpref::ArrayReferenceMetaPosition, elementbytes::Int)
   temp = gensym!(ls, "intermediateconstref")
   vloadcall = Expr(:call, lv(:_vload), mpref.mref.ptr)
   nindices = length(getindices(op))
@@ -117,6 +115,12 @@ function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementby
   pushpreamble!(ls, Expr(:(=), temp, vloadcall))
   pushpreamble!(ls, Expr(:(=), name(op), temp))
   pushpreamble!(ls, op, temp)
+  return temp
+end
+function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementbytes::Int)
+  op = Operation(length(operations(ls)), varname(mpref), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS, mpref.mref)
+  add_vptr!(ls, op)
+  temp = add_constant_vload!(ls, op, mpref, elementbytes)
   pushop!(ls, op, temp)
 end
 # This version has loop dependencies. var gets assigned to sym when lowering.

src/reconstruct_loopset.jl

@@ -716,8 +716,9 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
   # 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#")
   pushfirst!(ls.preamble.args, :(var"#lv#tuple#args#" = reassemble_tuple(Tuple{var"#LB#",var"#V#"}, var"#flattened#var#arguments#")))
+  post = hoist_constant_memory_accesses!(ls)
   # return @show avx_body(ls, var"#UNROLL#")
-  if last(var"#UNROLL#") > 1
+  q = if last(var"#UNROLL#") > 1
     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
     avx_threads_expr(
@@ -729,5 +730,6 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
     # return @show avx_body(ls, var"#UNROLL#")
     avx_body(ls, var"#UNROLL#")
   end
+  post === ls.preamble ? q : Expr(:block, q, post)
   # @show var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#"
 end

src/transforms.jl

@@ -0,0 +1,62 @@
+# file for misc loopset transforms
+
+function hoist_constant_memory_accesses!(ls::LoopSet)
+  hoist_stores = false
+  for op ∈ operations(ls)
+    if isload(op)
+      length(getindicesonly(op)) == 0 && host_constant_load!(ls, op)
+    elseif isstore(op) && iszero(length(getindicesonly(op)))
+      hoist_stores = true
+    end
+  end
+  hoist_stores && return hoist_constant_memory_accesses_nocheck!(ls)
+  ls.preamble
+end
+
+function hoist_constant_memory_accesses_nocheck!(ls::LoopSet)
+  post = Expr(:block)
+  for op ∈ operations(ls)
+    if isstore(op) && length(getindicesonly(op)) == 0
+      hoist_constant_store!(post, ls, op)
+    end
+  end
+  post
+end
+function hoist_constant_vload!(ls::LoopSet, op::Operation)
+  op.instr = LOOPCONSTANT
+  op.node_type = constant
+  add_constant_vload!(ls, op, ArrayReferenceMetaPosition(op.ref, parents(op), loopdependencies(op), reduceddependencies(op), name(op)), elementbytes)
+end
+
+function return_empty_reductinit(op::Operation, var::Symbol)
+  for (i,opp) ∈ enumerate(parents(op))
+    if (name(opp) === var) && (length(reduceddependencies(opp)) == 0) && (length(loopdependencies(opp)) == 0) && (length(children(opp)) == 1)
+      return opp
+    end
+    opcheck = return_empty_reductinit(opp, var)
+    opcheck === opp || return opcheck
+  end
+  return op
+end
+
+
+function hoist_constant_store!(q::Expr, ls::LoopSet, op::Operation)
+  op.instruction = DROPPEDCONSTANT
+  op.node_type = constant
+
+  opr = only(parents(op))
+  while opr.instruction.instr === :identity
+    opr.instruction = DROPPEDCONSTANT
+    opr.node_type = constant
+    opr = only(parents(opr))
+  end
+  push!(ls.outer_reductions, identifier(opr))
+  
+  init = return_empty_reductinit(opr, name(opr)).instruction.instr
+  pushpreamble!(ls, Expr(:(=), outer_reduct_init_typename(opr), Expr(:call, lv(:typeof), init)))
+  q = Expr(:block)
+  push!(q.args, Expr(:call, lv(:unsafe_store!), Expr(:call, lv(:pointer), op.ref.ptr), outer_reduction_to_scalar_reduceq!(q, opr, init)))
+  length(q.args) == 0 || pushpreamble!(ls, q) # creating `Expr` and pushing because `outer_reduction_to_scalar_reduceq!` uses `pushfirst!(q.args`, and we don't want it at the start of the preamble
+  return nothing
+end
+