More tweaks to correctly identifying family structure.

Author: chriselrod

src/LoopVectorization.jl

@@ -8,7 +8,7 @@ using VectorizationBase: REGISTER_SIZE, REGISTER_COUNT, extract_data, num_vector
     AbstractColumnMajorStridedPointer, AbstractRowMajorStridedPointer, AbstractSparseStridedPointer, AbstractStaticStridedPointer,
     PackedStridedPointer, SparseStridedPointer, RowMajorStridedPointer, StaticStridedPointer, StaticStridedStruct,
     maybestaticfirst, maybestaticlast, scalar_less, scalar_greater
-using SIMDPirates: VECTOR_SYMBOLS, evadd, evsub, evmul, evfdiv, vrange, reduced_add, reduced_prod, reduce_to_add, reduce_to_prod,
+using SIMDPirates: VECTOR_SYMBOLS, evadd, evsub, evmul, evfdiv, vrange, reduced_add, reduced_prod, reduce_to_add, reduce_to_prod, 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

src/add_compute.jl

@@ -48,25 +48,30 @@ function pushparent!(mpref::ArrayReferenceMetaPosition, parent::Operation)
     pushparent!(mpref.parents, mpref.loopdependencies, mpref.reduceddeps, parent)
 end
 function add_parent!(
-    parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var, elementbytes::Int, position::Int
+    vparents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var, elementbytes::Int, position::Int
 )
     parent = if var isa Symbol
-        getop(ls, var, elementbytes)
+        opp = getop(ls, var, elementbytes)
+        if iscompute(opp) && instruction(opp).instr === :identity && length(loopdependencies(opp)) < position && isone(length(parents(opp))) && name(opp) === name(first(parents(opp)))
+            first(parents(opp))
+        else
+            opp
+        end
     elseif var isa Expr #CSE candidate
         add_operation!(ls, gensym(:temporary), var, elementbytes, position)
     else # assumed constant
         add_constant!(ls, var, elementbytes)
         # add_constant!(ls, var, deps, gensym(:loopredefconst), elementbytes)
     end
-    pushparent!(parents, deps, reduceddeps, parent)
-end
-function add_reduction!(
-    parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var::Symbol, elementbytes::Int
-)
-    get!(ls.opdict, var) do
-        add_constant!(ls, var, elementbytes)
-    end
+    pushparent!(vparents, deps, reduceddeps, parent)
 end
+# function add_reduction!(
+#     vparents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var::Symbol, elementbytes::Int
+# )
+#     get!(ls.opdict, var) do
+#         add_constant!(ls, var, elementbytes)
+#     end
+# end
 function search_tree(opv::Vector{Operation}, var::Symbol) # relies on cycles being forbidden
     for opp ∈ opv
         name(opp) === var && return true
@@ -118,13 +123,6 @@ function add_reduced_deps!(op::Operation, reduceddeps::Vector{Symbol})
     nothing
 end
 
-# function substitute_op_in_parents!(
-#     vparents::Vector{Operation}, replacer::Operation, replacee::Operation, reduceddeps::Vector{Symbol}
-# )
-#     @show replacer replacee
-#     # 
-#     substitute_op_in_parents_recurse!(vparents, replacer, replacee)
-# end
 function substitute_op_in_parents!(
     vparents::Vector{Operation}, replacer::Operation, replacee::Operation, reduceddeps::Vector{Symbol}
 )
@@ -188,7 +186,7 @@ function add_reduction_update_parent!(
         if instr.instr ∈ (:-, :vsub!, :vsub, :/, :vfdiv!, :vfidiv!)
             update_deps!(deps, reduceddeps, reductinit)#parent) # deps and reduced deps will not be disjoint
         end
-    elseif !isouterreduction
+    elseif !isouterreduction && reductinit !== parent
         substitute_op_in_parents!(vparents, reductinit, parent, reduceddeps)
     end
     update_reduction_status!(vparents, reduceddeps, name(reductinit))
@@ -228,7 +226,8 @@ function add_compute!(
     for (ind,arg) ∈ enumerate(args)
         if var === arg
             reduction_ind = ind
-            add_reduction!(vparents, deps, reduceddeps, ls, arg, elementbytes)
+            # add_reduction!(vparents, deps, reduceddeps, ls, arg, elementbytes)
+            getop(ls, arg, elementbytes)
         elseif arg isa Expr
             isref, argref = tryrefconvert(ls, arg, elementbytes, varname(mpref))
             if isref
@@ -270,7 +269,8 @@ function add_compute!(
         parent = ls.opdict[var]
         setdiffv!(reduceddeps, deps, loopdependencies(parent))
         # parent = getop(ls, var, elementbytes)
-        if length(reduceddeps) == 0
+        # if length(reduceddeps) == 0
+        if all(!in(deps), reduceddeps)
             insert!(vparents, reduction_ind, parent)
             mergesetv!(deps, loopdependencies(parent))
             op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, vparents)

src/costs.jl

@@ -138,6 +138,8 @@ const COST = Dict{Symbol,InstructionCost}(
     :evsub => InstructionCost(4,0.5),
     :evmul => InstructionCost(4,0.5),
     :evfdiv => InstructionCost(13,4.0,-2.0),
+    :vsum => InstructionCost(6,2.0),
+    :vprod => InstructionCost(6,2.0),
     :reduced_add => InstructionCost(4,0.5),# ignoring reduction part of cost, might be nop
     :reduced_prod => InstructionCost(4,0.5),# ignoring reduction part of cost, might be nop
     :reduce_to_add => InstructionCost(0,0.0,0.0,0),
@@ -307,6 +309,10 @@ reduction_to_single_vector(x) = reduction_to_single_vector(reduction_instruction
 #     x == 1.0 ? :vsum : x == 2.0 ? :vprod : x == 5.0 ? :maximum : x == 6.0 ? :minimum : throw("Reduction not found.")
 # end
 # reduction_to_scalar(x) = reduction_to_scalar(reduction_instruction_class(x))
+function reduction_to_scalar(x::Float64)
+    x == ADDITIVE_IN_REDUCTIONS ? :vsum : x == MULTIPLICATIVE_IN_REDUCTIONS ? :vprod : x == MAX ? :vmaximum : x == MIN ? :vminimum : throw("Reduction not found.")
+end
+reduction_to_scalar(x) = reduction_to_scalar(reduction_instruction_class(x))
 function reduction_scalar_combine(x::Float64)
     # x == 1.0 ? :reduced_add : x == 2.0 ? :reduced_prod : x == 3.0 ? :reduced_any : x == 4.0 ? :reduced_all : x == 5.0 ? :reduced_max : x == 6.0 ? :reduced_min : throw("Reduction not found.")
     x == ADDITIVE_IN_REDUCTIONS ? :reduced_add : x == MULTIPLICATIVE_IN_REDUCTIONS ? :reduced_prod : x == MAX ? :reduced_max : x == MIN ? :reduced_min : throw("Reduction not found.")

src/lower_compute.jl

@@ -36,6 +36,39 @@ function parent_unroll_status(op::Operation, u₁loop::Symbol, u₂loop::Symbol,
     parents_u₁syms, parents_u₂syms
 end
 
+# """
+#     Requires a parents_op argument, because it may `===` parents(op), due to previous divergence, e.g. to handle unrolling.
+# """
+# function isreducingidentity!(q::Expr, op::Operation, parents_op::Vector{Operation}, U::Int, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, suffix)
+#     vparents = copy(parents_op) # don't mutate the original!
+#     for (i,opp) ∈ enumerate(parents_op)
+#         @show opp vectorized ∈ loopdependencies(opp), vectorized ∈ reducedchildren(opp) # must reduce
+#         @show vectorized, loopdependencies(opp), reducedchildren(opp) # must reduce
+#         if vectorized ∈ loopdependencies(opp) || vectorized ∈ reducedchildren(opp) # must reduce
+#             loopdeps = [l for l ∈ loopdependencies(opp) if l !== vectorized]
+#             @show opp
+#             reductinstruct = reduction_to_scalar(instruction(opp))
+            
+#             reducedparent = Operation(
+#                 opp.identifier, gensym(opp.variable), opp.elementbytes, Instruction(:LoopVectorization, reductinstruct), opp.node_type,
+#                 loopdeps, opp.reduced_deps, opp.parents, opp.ref, opp.reduced_children
+#             )
+#             pname,   pu₁,  pu₂ = variable_name_and_unrolled(opp, u₁loop, u₂loop, suffix)
+#             rpname, rpu₁, rpu₂ = variable_name_and_unrolled(reducedparent, u₁loop, u₂loop, suffix)
+#             @assert pu₁ == rpu₁ && pu₂ == rpu₂
+#             if rpu₁
+#                 for u ∈ 0:U-1
+#                     push!(q.args, Expr(:(=), Symbol(rpname,u), Expr(:call, lv(reductinstruct), Symbol(pname,u))))
+#                 end
+#             else
+#                 push!(q.args, Expr(:(=), rpname, Expr(:call, lv(reductinstruct), pname)))
+#             end
+#             vparents[i] = reducedparent
+#         end
+#     end
+#     vparents
+# end
+
 function lower_compute!(
     q::Expr, op::Operation, vectorized::Symbol, u₁loop::Symbol, u₂loop::Symbol, U::Int,
     suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing,
@@ -65,7 +98,7 @@ function lower_compute!(
     # end
     # unrollsym = isunrolled_sym(op, unrolled)
     if !opunrolled && any(parents_u₁syms) # TODO: Clean up this mess, refactor the naming code, putting it in one place and have everywhere else use it for easy equivalence.
-        parents_op = copy(parents_op)
+        parents_op = copy(parents_op) # don't mutate the original!
         for i ∈ eachindex(parents_u₁syms)
             parents_u₁syms[i] || continue
             parents_u₁syms[i] = false
@@ -115,7 +148,13 @@ function lower_compute!(
         end
     end
     # @show instr.instr
-    maskreduct = mask !== nothing && isreduct && vectorized ∈ reduceddependencies(op) #any(opp -> opp.variable === var, parents_op)
+    reduceddeps = reduceddependencies(op)
+    vecinreduceddeps = isreduct && vectorized ∈ reduceddeps
+    maskreduct = mask !== nothing && vecinreduceddeps #any(opp -> opp.variable === var, parents_op)
+    # if vecinreduceddeps && vectorized ∉ loopdependencies(op) # screen parent opps for those needing a reduction to scalar
+    #     # parents_op = reduce_vectorized_parents!(q, op, parents_op, U, u₁loop, u₂loop, vectorized, suffix)
+    #     isreducingidentity!(q, op, parents_op, U, u₁loop, u₂loop, vectorized, suffix) && return
+    # end    
     # if a parent is not unrolled, the compiler should handle broadcasting CSE.
     # because unrolled/tiled parents result in an unrolled/tiled dependendency,
     # we handle both the tiled and untiled case here.

src/lowering.jl

@@ -424,7 +424,7 @@ function lower(ls::LoopSet, us::UnrollSpecification)
 end
 
 function lower(ls::LoopSet, order, u₁loop, u₂loop, vectorized, u₁, u₂)
-    fillorder!(ls, order, u₁loop, u₂loop, u₂ != -1)
+    fillorder!(ls, order, u₁loop, u₂loop, u₂ != -1, vectorized)
     q = lower(ls, UnrollSpecification(ls, u₁loop, u₂loop, vectorized, u₁, u₂))
     iszero(length(ls.opdict)) && pushfirst!(q.args, Expr(:meta, :inline))
     q

src/operation_evaluation_order.jl

@@ -15,13 +15,31 @@
 # function iterate()
 
 # end
-function isnopidentity(op::Operation, u₁loop::Symbol, u₂loop::Symbol, suffix)
+
+function dependent_outer_reducts(ls::LoopSet, op)
+    for i ∈ ls.outer_reductions
+        search_tree(parents(operations(ls)[i]), name(op)) && return true
+    end
+    false
+end
+
+function isnopidentity(ls::LoopSet, op::Operation, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, suffix)
     parents_op = parents(op)
     if iscompute(op) && instruction(op).instr === :identity && name(first(parents_op)) === name(op) && isone(length(parents_op))
         mvar, u₁unrolledsym, u₂unrolledsym = variable_name_and_unrolled(op, u₁loop, u₂loop, suffix)
         parents_u₁syms, parents_u₂syms = parent_unroll_status(op, u₁loop, u₂loop, suffix)
         if (u₁unrolledsym == first(parents_u₁syms)) && ((!isnothing(suffix)) == parents_u₂syms[1])
-            true
+            #TODO: identifer(first(parents_op)) ∉ ls.outer_reductions is going to miss a lot of cases
+            #Should probably replace that with `DVec` (demoting Vec) types, that demote to scalar.
+            if (vectorized ∈ loopdependencies(first(parents_op)) && vectorized ∉ loopdependencies(op)) && !dependent_outer_reducts(ls, op)
+                op.instruction = reduction_to_scalar(instruction(first(parents_op)))
+                op.mangledvariable = gensym(op.mangledvariable)
+                false
+            else
+                true
+            end
+        else
+            false
         end
     else
         false
@@ -30,7 +48,6 @@ end
 
 function set_upstream_family!(adal::Vector{T}, op::Operation, val::T, ld::Vector{Symbol}, id::Int) where {T}
     adal[identifier(op)] == val && return # must already have been set
-    # @show op
     if ld != loopdependencies(op) || id == identifier(op)
         (adal[identifier(op)] = val)
     end
@@ -41,26 +58,27 @@ function set_upstream_family!(adal::Vector{T}, op::Operation, val::T, ld::Vector
 end
 
 function addoptoorder!(
-    lo::LoopOrder, included_vars::Vector{Bool}, place_after_loop::Vector{Bool}, op::Operation, loopsym::Symbol, _n::Int, u₁loop::Symbol, u₂loop::Symbol, loopistiled::Bool
+    ls::LoopSet, included_vars::Vector{Bool}, place_after_loop::Vector{Bool}, op::Operation,
+    loopsym::Symbol, _n::Int, u₁loop::Symbol, u₂loop::Symbol, vectorized::Symbol, loopistiled::Bool
 )
+    lo = ls.loop_order
     id = identifier(op)
     included_vars[id] && return nothing
     loopsym ∈ loopdependencies(op) || return nothing
     for opp ∈ parents(op) # ensure parents are added first
-        addoptoorder!(lo, included_vars, place_after_loop, opp, loopsym, _n, u₁loop, u₂loop, loopistiled)
+        addoptoorder!(ls, included_vars, place_after_loop, opp, loopsym, _n, u₁loop, u₂loop, vectorized, loopistiled)
     end
     included_vars[id] && return nothing
     included_vars[id] = true
     isunrolled = (u₁loop ∈ loopdependencies(op)) + 1
     istiled = u₂loop ∈ loopdependencies(op)
     # optype = Int(op.node_type) + 1
     after_loop = place_after_loop[id] + 1
-    # @show place_after_loop[id], op
     if !isloopvalue(op)
         if istiled
-            isnopidentity(op, u₁loop, u₂loop, 0) || push!(lo[isunrolled,2,after_loop,_n], op)
+            isnopidentity(ls, op, u₁loop, u₂loop, vectorized, 0) || push!(lo[isunrolled,2,after_loop,_n], op)
         else
-            isnopidentity(op, u₁loop, u₂loop, nothing) || push!(lo[isunrolled,1,after_loop,_n], op)
+            isnopidentity(ls, op, u₁loop, u₂loop, vectorized, nothing) || push!(lo[isunrolled,1,after_loop,_n], op)
         end
     end
     # isloopvalue(op) || push!(lo[isunrolled,istiled,after_loop,_n], op)
@@ -69,7 +87,7 @@ function addoptoorder!(
     nothing
 end
 
-function fillorder!(ls::LoopSet, order::Vector{Symbol}, unrolled::Symbol, tiled::Symbol, loopistiled::Bool)
+function fillorder!(ls::LoopSet, order::Vector{Symbol}, u₁loop::Symbol, u₂loop::Symbol, loopistiled::Bool, vectorized::Symbol)
     lo = ls.loop_order
     resize!(lo, length(ls.loopsymbols))
     ro = lo.loopnames # reverse order; will have same order as lo
@@ -85,7 +103,7 @@ function fillorder!(ls::LoopSet, order::Vector{Symbol}, unrolled::Symbol, tiled:
         ro[_n] = loopsym = order[n]
         #loopsym = order[n]
         for op ∈ ops
-            addoptoorder!( lo, included_vars, place_after_loop, op, loopsym, _n, unrolled, tiled, loopistiled )
+            addoptoorder!( ls, included_vars, place_after_loop, op, loopsym, _n, u₁loop, u₂loop, vectorized, loopistiled )
         end
     end
 end

test/gemm.jl

@@ -1,5 +1,5 @@
 @testset "GEMM" begin
-# using LoopVectorization, LinearAlgebra, Test; T = Float64
+ # using LoopVectorization, LinearAlgebra, Test; T = Float64
     Unum, Tnum = LoopVectorization.VectorizationBase.REGISTER_COUNT == 16 ? (3, 4) : (5, 5)
     AmulBtq1 = :(for m ∈ 1:size(A,1), n ∈ 1:size(B,2)
                  C[m,n] = zeroB
@@ -102,8 +102,17 @@
                  ΔCₘₙ += A[m,k] * B[k,n]
                  end
                  C[m,n] += ΔCₘₙ * factor
-                 end)
+                 end);
     lsAmuladd = LoopVectorization.LoopSet(Amuladdq);
+    Atmuladdq = :(for m ∈ 1:size(A,2), n ∈ 1:size(B,2)
+                 ΔCₘₙ = zero(eltype(C))
+                 for k ∈ 1:size(A,1)
+                 ΔCₘₙ += A[k,m] * B[k,n]
+                 end
+                 C[m,n] += ΔCₘₙ * factor
+                 end);
+    lsAtmuladd = LoopVectorization.LoopSet(Atmuladdq);
+    LoopVectorization.lower(lsAtmuladd, 2, 2)
     # lsAmuladd.operations
     # LoopVectorization.loopdependencies.(lsAmuladd.operations)
     # LoopVectorization.reduceddependencies.(lsAmuladd.operations)
@@ -520,7 +529,7 @@
             C[m,n] = Cmn_hi
         end
     end
-
+    
     function threegemms!(Ab, Bb, Cb, A, B, C)
         M, N = size(Cb); K = size(B,1)
         @avx for m in 1:M, k in 1:K, n in 1:N