Slightly more robust handling of unrolling in inner vectorized reductions

Author: chriselrod

Project.toml

@@ -28,7 +28,7 @@ SLEEFPirates = "0.6.14"
 Static = "0.2"
 ThreadingUtilities = "0.4.1"
 UnPack = "1"
-VectorizationBase = "0.19.22"
+VectorizationBase = "0.19.28"
 julia = "1.5"
 
 [extras]

src/LoopVectorization.jl

@@ -34,7 +34,8 @@ using Base.Meta: isexpr
 using DocStringExtensions
 import LinearAlgebra # for check_args
 
-using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast, rem_fast, max_fast, min_fast, log_fast, log2_fast, log10_fast
+using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast, rem_fast, max_fast, min_fast
+using SLEEFPirates: log_fast, log2_fast, log10_fast
 
 
 using ArrayInterface

src/broadcast.jl

@@ -368,7 +368,7 @@ function add_broadcast!(
         mergesetdiffv!(deps, loopdependencies(parent), reduceddependencies(parent))
     end
     op = Operation(
-        length(operations(ls)), destname, elementbytes, instr, compute, deps, NOPARENTS, parents
+        length(operations(ls)), destname, elementbytes, instr, compute, deps, NODEPENDENCY, parents
     )
     pushop!(ls, op, destname)
 end

src/codegen/lower_compute.jl

@@ -274,7 +274,7 @@ function parent_op_name(
     opp = parents_op[n]
     parent = mangledvar(opp)
     u = 0
-    if n == tiledouterreduction
+    if n == tiledouterreduction# && isvectorized(opp)
         parent = Symbol(parent, modsuffix)
     else
         if parents_u₂syms[n]
@@ -452,7 +452,11 @@ function lower_compute!(
         #     modsuffix = ls.unrollspecification[].u₁#getu₁full(ls, u₁)#u₁
         #     Symbol(mangledvar(op), '_', modsuffix)
         # else
-        modsuffix = 0#suffix % tiled_outerreduct_unroll(ls)
+        if u₁unrolledsym
+            modsuffix = 0
+        else
+            modsuffix = suffix % ls.ureduct[]
+        end
         Symbol(mangledvar(op), modsuffix)
         # end
         # dopartialmap = u₁ > 1
@@ -483,7 +487,8 @@ function lower_compute!(
             add_loopvalue!(instrcall, loopval, ua, u₁)
         elseif name(opp) === name(op)
             selfdep = n
-            if ((isvectorized(opp) && !isvectorized(op)) && !dependent_outer_reducts(ls, op)) ||
+            # @show mangledvar(op), name(opp), name(op)
+            if ((isvectorized(opp) && !isvectorized(op))) ||
                 (parents_u₁syms[n] != u₁unrolledsym) || (parents_u₂syms[n] != u₂unrolledsym)
 
                 selfopname, uₚ = parent_op_name(ls, parents_op, n, modsuffix, suffix_, parents_u₁syms, parents_u₂syms, u₁, opisvectorized, tiledouterreduction)
@@ -495,14 +500,26 @@ function lower_compute!(
             else
                 push!(instrcall.args, varsym)
             end
-        elseif ((!isu₂unrolled(op)) & isu₂unrolled(opp)) && (isouterreduction(ls, opp) != -1)
+        elseif ((!isu₂unrolled(op)) & isu₂unrolled(opp)) && (parents_u₂syms[n] & (!u₂unrolledsym))
+        # elseif parents_u₂syms[n] & (!u₂unrolledsym)
+            #&& (isouterreduction(ls, opp) != -1)
             # this checks if the parent is u₂ unrolled but this operation is not, in which case we need to reduce it.
+            # @show op opp
             reduced_u₂ = reduce_expr_u₂(mangledvar(opp), instruction(opp), ureduct(ls))
-            reduced_u₂ = reduce_parent!(q, ls, op, opp, reduced_u₂)
+            reducedparentname = gensym!(ls, "reducedop")
+            push!(q.args, Expr(:(=), reducedparentname, reduced_u₂))
+            reduced_u₂ = reduce_parent!(q, ls, op, opp, reducedparentname)
             push!(instrcall.args, reduced_u₂)
         else
             parent, uₚ = parent_op_name(ls, parents_op, n, modsuffix, suffix_, parents_u₁syms, parents_u₂syms, u₁, opisvectorized, tiledouterreduction)
             parent = reduce_parent!(q, ls, op, opp, parent)
+            # if instr.instr === :vfmadd_fast && tiledouterreduction > 0
+            #     @show mvar, varsym, selfopname
+            # end
+            # @show opp
+            # if instr.instr === :identity
+            #     @show isvectorized(op) isvectorized(opp)
+            # end
             if (selfdep == 0) && search_tree(parents(opp), name(op))
                 selfdep = n
                 push!(instrcall.args, parent)

src/codegen/lower_store.jl

@@ -181,7 +181,7 @@ function lower_tiled_store!(
 end
 
 function donot_tile_store(ls::LoopSet, op::Operation, vloop::Loop, reductfunc::Symbol, u₂::Int)
-    (!((reductfunc === Symbol("")) && all(op.ref.loopedindex))) || (u₂ ≤ 1) || isconditionalmemop(op) && return true
+    ((!((reductfunc === Symbol("")) && all(op.ref.loopedindex))) || (u₂ ≤ 1) || isconditionalmemop(op)) && return true
     rejectcurly(op) && return true
     omop = offsetloadcollection(ls)
     batchid, opind = omop.batchedcollectionmap[identifier(op)]

src/codegen/lowering.jl

@@ -441,10 +441,9 @@ function initialize_outer_reductions!(
     @unpack u₁, u₂ = us
     Umax = u₂ == -1 ? _Umax : u₁
     reduct_zero = reduction_zero(op.instruction)
-    isvectorized = vectorized ∈ reduceddependencies(op)
     typeTr = ELTYPESYMBOL
     u₁u, u₂u = isunrolled_sym(op, getloop(ls, us.u₁loopnum).itersymbol, getloop(ls, us.u₂loopnum).itersymbol, getloop(ls, us.vloopnum).itersymbol)#, u₂)
-    z = if isvectorized
+    z = if isvectorized(op)
         if Umax == 1 || !u₁u
             if reduct_zero === :zero
                 Expr(:call, lv(:_vzero), VECTORWIDTHSYMBOL, typeTr, rs)
@@ -792,25 +791,15 @@ function calc_Ureduct!(ls::LoopSet, us::UnrollSpecification)
             if u₁ui == -1
                 u₁ui = Int(u₁u)
                 u₂ui = Int(u₁u)
-            else
-                @assert (u₁ui == Int(u₁u)) & (u₂ui == Int(u₁u)) "Doesn't currenly andle differently unrolled reductions yet, please file an issue with an example."
+            elseif !((u₁ui == Int(u₁u)) & (u₂ui == Int(u₁u)))
+                throw(ArgumentError("Doesn't currenly handle differently unrolled reductions yet, please file an issue with an example."))
             end
         end
         if u₁ui % Bool
             u₁
-        #     push!(q.args, Expr(:(=), Symbol(mvar, '_', u₁), z))
         else
             u₂
-        #     for u ∈ 0:_Umax-1
-        #         # push!(q.args, Expr(:(=), Symbol(mvar, '_', u), z))
-        #         push!(q.args, Expr(:(=), Symbol(mvar, u), z))
-        #     end
         end
-    #     u₁loopnum == vloopnum
-    #     u₂
-    # else
-    #     u₁
-        # u₁#tiled_outerreduct_unroll(us)
     end
     ls.ureduct[] = ur
 end
@@ -935,6 +924,8 @@ function isunrolled_sym(op::Operation, u₁loop::Symbol, u₂loop::Symbol, vloop
     ((u₂max > 1) | accesses_memory(op)) ? isunrolled_sym(op, u₁loop, u₂loop, vloop) : (isunrolled_sym(op, u₁loop), false)
 end
 
+
+
 function variable_name(op::Operation, suffix::Int)
     mvar = mangledvar(op)
     suffix == -1 ? mvar : Symbol(mvar, suffix, :_)

src/codegen/operation_evaluation_order.jl

@@ -34,17 +34,12 @@ function isnopidentity(ls::LoopSet, op::Operation, u₁loop::Symbol, u₂loop::S
         # if (u₁unrolledsym == first(parents_u₁syms)) && (isu₂unrolled(op) == parents_u₂syms[1])
         opp = only(parents_op)
         if (isu₁unrolled(op) == isu₁unrolled(opp)) & (isu₂unrolled(op) == isu₂unrolled(opp))
-            #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.
-            #TODO: document (after finding out...) why only checking `isvectorized(first(parents_op))` -- why not `any(isvectorized, parents_op)`???
-            if (isvectorized(opp) && !isvectorized(op)) && !dependent_outer_reducts(ls, op)
+            true
+        else
+            if isvectorized(opp) & (!isvectorized(op))
                 op.instruction = reduction_to_scalar(instruction(opp))
                 op.mangledvariable = gensym(op.mangledvariable)
-                false
-            else
-                true
             end
-        else
             false
         end
     else

src/modeling/graphs.jl

@@ -592,7 +592,7 @@ function cacheunrolled!(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, vloopsy
         end
     end
     for op ∈ operations(ls)
-        setunrolled!(op, u₁loop, u₂loop, vloopsym)
+        setunrolled!(ls, op, u₁loop, u₂loop, vloopsym)
         if accesses_memory(op)
             rc = rejectcurly(ls, op, u₁loop, vloopsym)
             op.rejectcurly = rc
@@ -606,6 +606,47 @@ function cacheunrolled!(ls::LoopSet, u₁loop::Symbol, u₂loop::Symbol, vloopsy
         end
     end
 end
+function setunrolled!(ls::LoopSet, op::Operation, u₁loopsym, u₂loopsym, vectorized)
+    op.u₁unrolled =  u₁loopsym ∈ loopdependencies(op)
+    op.u₂unrolled =  u₂loopsym ∈ loopdependencies(op)
+    op.vectorized = vectorized ∈ loopdependencies(op)
+    if isconstant(op)
+        u₁ = op.u₁unrolled
+        u₂ = op.u₂unrolled
+        v  = op.vectorized
+        for opp ∈ children(op)
+            u₁ = u₁ &&  u₁loopsym ∈ loopdependencies(opp)
+            u₂ = u₂ &&  u₂loopsym ∈ loopdependencies(opp)
+            v  = v  && vectorized ∈ loopdependencies(opp)
+        end
+        if isouterreduction(ls, op) ≠ -1 && !all((u₁,u₂,v))
+            opv = true
+            for opp ∈ parents(op)
+                if iscompute(opp) && instruction(opp).instr ≢ :identity
+                    opv = false
+                    break
+                end
+            end
+            if opv
+                if !u₁ && u₁loopsym ∈ reduceddependencies(op)
+                    u₁ = true
+                end
+                if !u₂ && u₂loopsym ∈ reduceddependencies(op)
+                    u₂ = true
+                end
+                if !v && vectorized ∈ reduceddependencies(op)
+                    v = true
+                end
+            end
+        end
+        op.u₁unrolled = u₁
+        op.u₂unrolled = u₂
+        op.vectorized = v
+    end
+    # op.u₁unrolled, op.u₂unrolled = isunrolled_sym(op, u₁loopsym, u₂loopsym, vectorized)
+    nothing
+end
+
 rejectcurly(op::Operation) = op.rejectcurly
 rejectinterleave(op::Operation) = op.rejectinterleave
 num_loops(ls::LoopSet) = length(ls.loops)

src/modeling/operations.jl

@@ -200,21 +200,19 @@ mutable struct Operation <: AbstractLoopOperation
     rejectinterleave::Bool
     function Operation(
         identifier::Int,
-        variable,
-        elementbytes,
-        instruction,
-        node_type,
-        dependencies = Symbol[],
-        reduced_deps = Symbol[],
-        parents = Operation[],
+        variable::Symbol,
+        elementbytes::Int,
+        instruction::Union{Symbol,Instruction},
+        node_type::OperationType,
+        dependencies::Vector{Symbol} = Symbol[],
+        reduced_deps::Vector{Symbol} = Symbol[],
+        parents::Vector{Operation} = Operation[],
         ref::ArrayReferenceMeta = NOTAREFERENCE,
-        reduced_children = Symbol[]
+        reduced_children::Vector{Symbol} = Symbol[]
     )
         new(
             identifier, variable, elementbytes, instruction, node_type,
-            convert(Vector{Symbol},dependencies),
-            convert(Vector{Symbol},reduced_deps),
-            convert(Vector{Operation},parents), Operation[],
+            dependencies, reduced_deps, parents, Operation[],
             ref, Symbol("##", variable, :_),
             reduced_children
         )
@@ -224,19 +222,6 @@ end
 isu₁unrolled(op::Operation) = op.u₁unrolled
 isu₂unrolled(op::Operation) = op.u₂unrolled
 isvectorized(op::Operation) = op.vectorized
-function setunrolled!(op::Operation, u₁loopsym, u₂loopsym, vectorized)
-    op.u₁unrolled = u₁loopsym ∈ loopdependencies(op)
-    op.u₂unrolled = u₂loopsym ∈ loopdependencies(op)
-    op.vectorized = vectorized ∈ loopdependencies(op)
-    if isconstant(op)
-        for opp ∈ children(op)
-            op.u₁unrolled = op.u₁unrolled && u₁loopsym ∈ loopdependencies(opp)
-            op.u₂unrolled = op.u₂unrolled && u₂loopsym ∈ loopdependencies(opp)
-            op.vectorized = op.vectorized && vectorized ∈ loopdependencies(opp)            
-        end
-    end
-    nothing
-end
 
 function matches(op1::Operation, op2::Operation)
     op1 === op2 && return true

src/parse/add_stores.jl

@@ -123,7 +123,7 @@ function add_conditional_store!(ls::LoopSet, LHS, condop::Operation, storeop::Op
     #         return cse_store!(ls, op)
     #     end
     # end
-    op = Operation( id, name(mref), elementbytes, :conditionalstore!, memstore, ldref, NODEPENDENCY, storeparents, mref )
+    op = Operation( id, name(mref), elementbytes, :conditionalstore!, memstore, ldref, reduceddependencies(storeop), storeparents, mref )
     add_unique_store!(ls, op)
 end