Fix for prefetch-check in case of dependency on non-loop ops.

Author: chriselrod

src/LoopVectorization.jl

@@ -16,7 +16,7 @@ using SIMDPirates: VECTOR_SYMBOLS, evadd, evsub, evmul, evfdiv, vrange,
     reduced_add, reduced_prod, reduce_to_add, reduced_max, reduced_min, vsum, vprod, vmaximum, vminimum,
     sizeequivalentfloat, sizeequivalentint, vadd!, vsub!, vmul!, vfdiv!, vfmadd!, vfnmadd!, vfmsub!, vfnmsub!,
     vfmadd231, vfmsub231, vfnmadd231, vfnmsub231, sizeequivalentfloat, sizeequivalentint, #prefetch,
-    vmullog2, vmullog10, vdivlog2, vdivlog10, vmullog2add!, vmullog10add!, vdivlog2add!, vdivlog10add!, vfmaddaddone, vadd1
+    vmullog2, vmullog10, vdivlog2, vdivlog10, vmullog2add!, vmullog10add!, vdivlog2add!, vdivlog10add!, vfmaddaddone, vadd1, relu
 using SLEEFPirates: pow
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
 using LinearAlgebra: Adjoint, Transpose

src/add_stores.jl

@@ -57,15 +57,14 @@ function add_store!(
     mpref = array_reference_meta!(ls, array, rawindices, elementbytes, var)
     add_store!(ls, mpref, elementbytes)
 end
-function add_simple_store!(ls::LoopSet, parent::Operation, ref::ArrayReference, elementbytes::Int)
-    mref = ArrayReferenceMeta(
-        ref, fill(true, length(getindices(ref)))
-    )
-    op = Operation( ls, name(mref), elementbytes, :setindex!, memstore, getindices(ref), NODEPENDENCY, [parent], mref )
+function add_simple_store!(ls::LoopSet, parent::Operation, mref::ArrayReferenceMeta, elementbytes::Int)
+    op = Operation( ls, name(mref), elementbytes, :setindex!, memstore, getindices(mref.ref), NODEPENDENCY, [parent], mref )
     add_unique_store!(ls, op)
 end
-function add_simple_store!(ls::LoopSet, var::Symbol, ref::ArrayReference, elementbytes::Int)
-    add_simple_store!(ls, getop(ls, var, elementbytes), ref, elementbytes)
+function add_simple_store!(ls::LoopSet, var::Union{Symbol,Operation}, ref::Union{ArrayReference,ArrayReferenceMeta}, elementbytes::Int)
+    parent = isa(var, Symbol) ? getop(ls, var, elementbytes) : var
+    mref = isa(ref, ArrayReference) ? ArrayReferenceMeta(ref, fill(true, length(getindices(ref)))) : ref
+    add_simple_store!(ls, parent, mref, elementbytes)
 end
 function add_store_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int)
     array, raw_indices = ref_from_ref!(ls, ex)

src/costs.jl

@@ -183,6 +183,7 @@ const COST = Dict{Symbol,InstructionCost}(
     :iseven => InstructionCost(1, 0.5),
     :max => InstructionCost(4,0.5),
     :min => InstructionCost(4,0.5),
+    :relu => InstructionCost(4,0.5),
     # Instruction(:ifelse) => InstructionCost(1, 0.5),
     :vifelse => InstructionCost(1, 0.5),
     :inv => InstructionCost(13,4.0,-2.0,1),
@@ -452,6 +453,7 @@ const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     # typeof(SLEEFPirates.tanh_fast) => :tanh_fast,
     typeof(max) => :max,
     typeof(min) => :min,
+    typeof(relu) => :relu,
     typeof(<<) => :<<,
     typeof(>>) => :>>,
     typeof(>>>) => :>>>,

src/graphs.jl

@@ -70,6 +70,8 @@ function Loop(itersymbol::Symbol, start::Union{Int,Symbol}, stop::Union{Int,Symb
 end
 Base.length(loop::Loop) = 1 + loop.stophint - loop.starthint
 isstaticloop(loop::Loop) = loop.startexact & loop.stopexact
+
+
 function startloop(loop::Loop, itersymbol)
     startexact = loop.startexact
     if startexact
@@ -353,11 +355,12 @@ function oporder(ls::LoopSet)
 end
 names(ls::LoopSet) = ls.loop_order.loopnames
 reversenames(ls::LoopSet) = ls.loop_order.bestorder
-function getloopid(ls::LoopSet, s::Symbol)::Int
+function getloopid_or_nothing(ls::LoopSet, s::Symbol)
     for (loopnum,sym) ∈ enumerate(ls.loopsymbols)
         s === sym && return loopnum
     end
 end
+getloopid(ls::LoopSet, s::Symbol) = getloopid_or_nothing(ls, s)::Int
 getloop(ls::LoopSet, s::Symbol) = ls.loops[getloopid(ls, s)]
 # getloop(ls::LoopSet, i::Integer) = ls.loops[i]
 getloopsym(ls::LoopSet, i::Integer) = ls.loopsymbols[i]
@@ -366,6 +369,10 @@ Base.length(ls::LoopSet, s::Symbol) = length(getloop(ls, s))
 # isstaticloop(ls::LoopSet, s::Symbol) = isstaticloop(getloop(ls,s))
 # looprangehint(ls::LoopSet, s::Symbol) = length(getloop(ls, s))
 # looprangesym(ls::LoopSet, s::Symbol) = getloop(ls, s).rangesym
+
+"""
+getop only works while construction a LoopSet object. You cannot use it while lowering.
+"""
 getop(ls::LoopSet, var::Number, elementbytes) = add_constant!(ls, var, elementbytes)
 function getop(ls::LoopSet, var::Symbol, elementbytes::Int)
     get!(ls.opdict, var) do
@@ -379,6 +386,16 @@ function getop(ls::LoopSet, var::Symbol, deps, elementbytes::Int)
 end
 getop(ls::LoopSet, i::Int) = ls.operations[i]
 
+# """ 
+# Returns an operation with the same name as `s`.
+# """
+# function getoperation(ls::LoopSet, s::Symbol)
+#     for op ∈ Iterators.Reverse(operations(ls))
+#         name(op) === s && return op
+#     end
+#     throw("Symbol $s not found among operations(ls).")
+# end
+
 function Operation(
     ls::LoopSet, variable, elementbytes, instruction,
     node_type, dependencies, reduced_deps, parents, ref = NOTAREFERENCE
@@ -731,10 +748,39 @@ function UnrollSpecification(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, ve
 end
 
 """
+  looplengthprod(ls::LoopSet)
+
 Convert to `Float64` for the sake of non-64 bit platforms.
 """
 looplengthprod(ls::LoopSet) = prod(Float64 ∘ length, ls.loops)
 
+
+function looplength(ls::LoopSet, s::Symbol)
+    # search_tree(parents(operations(ls)[i]), name(op)) && return true
+    id = getloopid_or_nothing(ls, s)
+    if isnothing(id)
+        l = 0.0
+        # TODO: we could double count a loop.
+        for op ∈ operations(ls)
+            name(op) === s || continue
+            for opp ∈ parents(op)
+                if isloopvalue(opp)
+                    oppname = first(loopdependencies(opp))
+                    l += looplength(ls, oppname)
+                elseif iscompute(opp)
+                    oppname = name(opp)
+                    l += looplength(ls, oppname)
+                    # TODO elseif isconstant(opp)
+                end
+            end
+            l += 1 - length(parents(op))
+        end
+        l
+    else
+        Float64(length(ls.loops[id]))
+    end
+end
+
 # function getunrolled(ls::LoopSet)
 #     order = names(ls)
 #     us = ls.unrollspecification[]

src/lower_load.jl

@@ -87,7 +87,7 @@ function prefetchisagoodidea(ls::LoopSet, op::Operation, td::UnrollArgs)
             innermostloopindv = findall(map(isequal(innermostloopsym), getindices(op)))
             isone(length(innermostloopindv)) || return 0
             innermostloopind = first(innermostloopindv)
-            if prod(s -> Float64(length(getloop(ls, s))), @view(indices[1:innermostloopind-1])) ≥ 120.0 && length(getloop(ls, innermostloopsym)) ≥ 120
+            if prod(s -> Float64(looplength(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120.0 && length(getloop(ls, innermostloopsym)) ≥ 120
                 if op.ref.ref.offsets[innermostloopind] < 120
                     for opp ∈ operations(ls)
                         iscompute(opp) && (innermostloopsym ∈ loopdependencies(opp)) && load_constrained(opp, u₁loopsym, u₂loopsym, innermostloopsym, true) && return 0

src/lowering.jl

@@ -671,6 +671,7 @@ end
 function lower_unrollspec(ls::LoopSet)
     us = ls.unrollspecification[]
     @unpack vectorizedloopnum, u₁, u₂ = us
+    # @show u₁, u₂
     order = names(ls)
     vectorized = order[vectorizedloopnum]
     setup_preamble!(ls, us)

src/operations.jl

@@ -242,7 +242,7 @@ function Base.show(io::IO, op::Operation)
         ref = Expr(:ref, name(op.ref)); append!(ref.args, getindices(op))
         print(io, Expr(:(=), ref, name(first(parents(op)))))
     elseif isloopvalue(op)
-        print(io, Expr(:(=), op.variable, op.variable))
+        print(io, Expr(:(=), op.variable, first(loopdependencies(op))))
     end
 end
 

src/vectorizationbase_extensions.jl

@@ -23,6 +23,13 @@ end
         Boff.offsets
     )
 end
+
+@inline function VectorizationBase.stridedpointer_for_broadcast(A::OffsetArrays.OffsetArray)
+    OffsetStridedPointer(
+        VectorizationBase.stridedpointer_for_broadcast(parent(A)),
+        VectorizationBase.staticmul(T, VectorizationBase.filter_strides_by_dimequal1(Base.tail(size(A)), VectorizationBase.tailstrides(A)))
+    )
+end
 @inline function Base.transpose(A::OffsetStridedPointer)
     OffsetStridedPointer(
         transpose(A.ptr), A.offsets