Updated reduction to track through multiple operations.

Author: chriselrod

src/add_compute.jl

@@ -61,49 +61,66 @@ function add_reduction!(
     end
     # pushparent!(parents, deps, reduceddeps, parent)
 end
+function search_tree(opv::Vector{Operation}, var::Symbol) # relies on cycles being forbidden
+    for opp ∈ opv
+        name(opp) === var && return true
+        search_tree(parents(opp), var) && return true
+    end
+    false
+end
+function update_reduction_status!(parentvec::Vector{Operation}, deps::Vector{Symbol}, parent::Symbol)
+    for opp ∈ parentvec
+        if name(opp) === parent
+            mergesetv!(reducedchildren(opp), deps)
+            break
+        elseif search_tree(parents(opp), parent)
+            mergesetv!(reducedchildren(opp), deps)
+            update_reduction_status!(parents(opp), deps, parent)
+            break
+        end
+    end
+end
 function add_reduction_update_parent!(
     parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet,
-    var::Symbol, instr::Symbol, elementbytes::Int = 8
+    var::Symbol, instr::Symbol, directdependency::Bool, elementbytes::Int = 8
 )
     parent = getop(ls, var, elementbytes)
-    isloopconstant = parent.instruction === LOOPCONSTANT
+    isouterreduction = parent.instruction === LOOPCONSTANT
     Instr = Instruction(instr)
+    instrclass = reduction_instruction_class(Instr) # key allows for faster lookups
     # if parent is not an outer reduction...
-    if !isloopconstant
+    if !isouterreduction
         # and parent is not a reduction_zero
-        reduct_zero = REDUCTION_ZERO[Instr]
-        reductcombine::Symbol = @static VERSION < v"1.3" ? last(REDUCTION_SCALAR_COMBINE[Instr].args).value : REDUCTION_SCALAR_COMBINE[Instr].name
+        reduct_zero = reduction_zero(instrclass)
+        reductcombine = reduction_scalar_combine(instrclass)
         reductsym = gensym(:reduction)
         reductinit = add_constant!(ls, Expr(:call, reduct_zero, ls.T), loopdependencies(parent), reductsym, reduct_zero, elementbytes)
         if isconstant(parent) && reduct_zero === parent.instruction.mod #we can use parent op as initialization.
-            reductcombine = REDUCTION_COMBINETO[reductcombine]
-        # else # we cannot use parent op as initialization.
+            reductcombine = reduction_combine_to(instrclass)
         end
     else
         reductinit = parent
         reductsym = var
         reductcombine = Symbol("")
     end
-    # mergesetv!(reduceddeps, deps)
-    # if length(reduceddependencies(reductinit)) == 0
-        # setdiffv!(reduceddeps, deps, loopdependencies(reductinit))
-    # else
     setdiffv!(reduceddeps, deps, loopdependencies(reductinit))
-    # end
-    # mergesetv!(reduceddependencies(reductinit), reduceddeps)
-    pushparent!(parents, deps, reduceddeps, reductinit)#parent) # deps and reduced deps will not be disjoint
+    combineddeps = copy(deps); mergesetv!(combineddeps, reduceddeps)
+    directdependency && pushparent!(parents, deps, reduceddeps, reductinit)#parent) # deps and reduced deps will not be disjoint
+    update_reduction_status!(parents, combineddeps, name(reductinit))
+    # this is the op added by add_compute
     op = Operation(length(operations(ls)), reductsym, elementbytes, instr, compute, deps, reduceddeps, parents)
     parent.instruction === LOOPCONSTANT && push!(ls.outer_reductions, identifier(op))
     opout = pushop!(ls, op, var) # note this overwrites the entry in the operations dict, but not the vector
-    isloopconstant && return opout
-    # create child
+    isouterreduction && return opout
+    # create child op, which is the reduction combination
     childdeps = Symbol[]; childrdeps = Symbol[]; childparents = Operation[]
     pushparent!(childparents, childdeps, childrdeps, op) # reduce op
     pushparent!(childparents, childdeps, childrdeps, parent) # to
     child = Operation(
         length(operations(ls)), name(parent), elementbytes, reductcombine, compute, childdeps, childrdeps, childparents
     )
     pushop!(ls, child, name(parent))
+    opout
 end
 function add_compute!(
     ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8,
@@ -136,8 +153,8 @@ function add_compute!(
             add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes)
         end
     end
-    if reduction # arg[reduction] is the reduction
-        add_reduction_update_parent!(parents, deps, reduceddeps, ls, var, instr, elementbytes)
+    if reduction || search_tree(parents, var)
+        add_reduction_update_parent!(parents, deps, reduceddeps, ls, var, instr, reduction, elementbytes)
     else
         op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, parents)
         pushop!(ls, op, var)

src/condense_loopset.jl

@@ -44,6 +44,7 @@ struct OperationStruct
     instruction::Instruction
     loopdeps::UInt64
     reduceddeps::UInt64
+    childdeps::UInt64
     parents::UInt64
     node_type::OperationType
     array::UInt8
@@ -60,7 +61,8 @@ filled_4byte_chunks(u::UInt64) = 16 - (leading_zeros(u) >>> 2)
 filled_8byte_chunks(u::UInt64) = 8 - (leading_zeros(u) >>> 3)
 
 num_loop_deps(os::OperationStruct) = filled_4byte_chunks(os.loopdeps)
-num_reduced_deps(os::OperationStruct) = filled_4byte_chunks(os.reduced_deps)
+num_reduced_deps(os::OperationStruct) = filled_4byte_chunks(os.reduceddeps)
+num_child_deps(os::OperationStruct) = filled_4byte_chunks(os.childdeps)
 num_parents(os::OperationStruct) = filled_4byte_chunks(os.parents)
 
 function shifted_loopset(ls::LoopSet, loopsyms::Vector{Symbol})
@@ -73,6 +75,7 @@ function shifted_loopset(ls::LoopSet, loopsyms::Vector{Symbol})
 end
 loopdeps_uint(ls::LoopSet, op::Operation) = shifted_loopset(ls, loopdependencies(op))
 reduceddeps_uint(ls::LoopSet, op::Operation) = shifted_loopset(ls, reduceddependencies(op))
+childdeps_uint(ls::LoopSet, op::Operation) = shifted_loopset(ls, reducedchildren(op))
 function parents_uint(ls::LoopSet, op::Operation)
     p = zero(UInt64)
     for parent ∈ parents(op)
@@ -85,10 +88,11 @@ function OperationStruct(ls::LoopSet, op::Operation)
     instr = instruction(op)
     ld = loopdeps_uint(ls, op)
     rd = reduceddeps_uint(ls, op)
+    cd = childdeps_uint(ls, op)
     p = parents_uint(ls, op)
     array = accesses_memory(op) ? findmatchingarray(ls, vptr(op.ref)) : 0x00
     OperationStruct(
-        instr, ld, rd, p, op.node_type, array
+        instr, ld, rd, cd, p, op.node_type, array
     )
 end
 ## turn a LoopSet into a type object which can be used to reconstruct the LoopSet.

src/costs.jl

@@ -156,108 +156,58 @@ const COST = Dict{Instruction,InstructionCost}(
 #     COST[Symbol("typeof(", lower(k), ")")] = v
 # end
 
-const CORRESPONDING_REDUCTION = Dict{Instruction,Instruction}(
-    Instruction(:+) => Instruction(:vsum),
-    Instruction(:-) => Instruction(:vsum),
-    Instruction(:*) => Instruction(:vprod),
-    Instruction(:vadd) => Instruction(:vsum),
-    Instruction(:vsub) => Instruction(:vsum),
-    Instruction(:vmul) => Instruction(:vprod),
-    Instruction(:evadd) => Instruction(:vsum),
-    Instruction(:evsub) => Instruction(:vsum),
-    Instruction(:evmul) => Instruction(:vprod),
-    Instruction(:&) => Instruction(:vall),
-    Instruction(:|) => Instruction(:vany),
-    Instruction(:muladd) => Instruction(:vsum),
-    Instruction(:fma) => Instruction(:vsum),
-    Instruction(:vmuladd) => Instruction(:vsum),
-    Instruction(:vfma) => Instruction(:vsum),
-    Instruction(:vfmadd) => Instruction(:vsum),
-    Instruction(:vfmsub) => Instruction(:vsum),
-    Instruction(:vfnmadd) => Instruction(:vsum),
-    Instruction(:vfnmsub) => Instruction(:vsum),
-    Instruction(:vfmadd_fast) => Instruction(:vsum),
-    Instruction(:vfmsub_fast) => Instruction(:vsum),
-    Instruction(:vfnmadd_fast) => Instruction(:vsum),
-    Instruction(:vfnmsub_fast) => Instruction(:vsum)
-)
-const REDUCTION_TRANSLATION = Dict{Instruction,Instruction}(
-    Instruction(:+) => Instruction(:evadd),
-    Instruction(:vadd) => Instruction(:evadd),
-    Instruction(:*) => Instruction(:evmul),
-    Instruction(:vmul) => Instruction(:evmul),
-    Instruction(:-) => Instruction(:evadd),
-    Instruction(:vsub) => Instruction(:evadd),
-    Instruction(:/) => Instruction(:evmul),
-    Instruction(:vfdiv) => Instruction(:evmul),
-    Instruction(:muladd) => Instruction(:evadd),
-    Instruction(:fma) => Instruction(:evadd),
-    Instruction(:vmuladd) => Instruction(:evadd),
-    Instruction(:vfma) => Instruction(:evadd),
-    Instruction(:vfmadd) => Instruction(:evadd),
-    Instruction(:vfmsub) => Instruction(:evadd),
-    Instruction(:vfnmadd) => Instruction(:evadd),
-    Instruction(:vfnmsub) => Instruction(:evadd),
-    Instruction(:vfmadd_fast) => Instruction(:evadd),
-    Instruction(:vfmsub_fast) => Instruction(:evadd),
-    Instruction(:vfnmadd_fast) => Instruction(:evadd),
-    Instruction(:vfnmsub_fast) => Instruction(:evadd)
-)
-const REDUCTION_ZERO = Dict{Instruction,Symbol}(
-    Instruction(:+) => :zero,
-    Instruction(:vadd) => :zero,
-    Instruction(:evadd) => :zero,
-    Instruction(:*) => :one,
-    Instruction(:vmul) => :one,
-    Instruction(:evmul) => :one,
-    Instruction(:-) => :zero,
-    Instruction(:vsub) => :zero,
-    Instruction(:evsub) => :zero,
-    Instruction(:/) => :one,
-    Instruction(:vfdiv) => :one,
-    Instruction(:evfdiv) => :one,
-    Instruction(:muladd) => :zero,
-    Instruction(:fma) => :zero,
-    Instruction(:vmuladd) => :zero,
-    Instruction(:vfma) => :zero,
-    Instruction(:vfmadd) => :zero,
-    Instruction(:vfmsub) => :zero,
-    Instruction(:vfnmadd) => :zero,
-    Instruction(:vfnmsub) => :zero,
-    Instruction(:vfmadd_fast) => :zero,
-    Instruction(:vfmsub_fast) => :zero,
-    Instruction(:vfnmadd_fast) => :zero,
-    Instruction(:vfnmsub_fast) => :zero
-)
+const ADDITIVE_IN_REDUCTIONS = 1.0
+const MULTIPLICATIVE_IN_REDUCTIONS = 2.0
+const ANY = 3.0
+const ALL = 4.0
 
-const LVGETPROP = @static VERSION < v"1.3" ? Expr : GlobalRef
-# Fast functions, because common pattern is
-const REDUCTION_SCALAR_COMBINE = Dict{Instruction,LVGETPROP}(
-    Instruction(:+) => lv(:reduced_add),
-    Instruction(:vadd) => lv(:reduced_add),
-    Instruction(:*) => lv(:reduced_prod),
-    Instruction(:vmul) => lv(:reduced_prod),
-    Instruction(:-) => lv(:reduced_add),
-    Instruction(:vsub) => lv(:reduced_add),
-    Instruction(:/) => lv(:reduced_prod),
-    Instruction(:vfdiv) => lv(:reduced_prod),
-    Instruction(:muladd) => lv(:reduced_add),
-    Instruction(:fma) => lv(:reduced_add),
-    Instruction(:vmuladd) => lv(:reduced_add),
-    Instruction(:vfma) => lv(:reduced_add),
-    Instruction(:vfmadd) => lv(:reduced_add),
-    Instruction(:vfmsub) => lv(:reduced_add),
-    Instruction(:vfnmadd) => lv(:reduced_add),
-    Instruction(:vfnmsub) => lv(:reduced_add),
-    Instruction(:vfmadd_fast) => lv(:reduced_add),
-    Instruction(:vfmsub_fast) => lv(:reduced_add),
-    Instruction(:vfnmadd_fast) => lv(:reduced_add),
-    Instruction(:vfnmsub_fast) => lv(:reduced_add)
-)
-const REDUCTION_COMBINETO = Dict{Symbol,Symbol}(
-    :reduced_add => :reduce_to_add,
-    :reduced_prod => :reduce_to_prod
+const REDUCTION_CLASS = Dict{Instruction,Float64}(
+    Instruction(:+) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:-) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:*) => MULTIPLICATIVE_IN_REDUCTIONS,
+    Instruction(:vadd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vsub) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vmul) => MULTIPLICATIVE_IN_REDUCTIONS,
+    Instruction(:evadd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:evsub) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:evmul) => MULTIPLICATIVE_IN_REDUCTIONS,
+    Instruction(:&) => ALL,
+    Instruction(:|) => ANY,
+    Instruction(:muladd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:fma) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vmuladd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfma) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfmadd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfmsub) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfnmadd) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfnmsub) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfmadd_fast) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfmsub_fast) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfnmadd_fast) => ADDITIVE_IN_REDUCTIONS,
+    Instruction(:vfnmsub_fast) => ADDITIVE_IN_REDUCTIONS
 )
+reduction_instruction_class(instr::Symbol) = get(REDUCTION_CLASS, Instruction(instr), NaN)
+reduction_instruction_class(instr::Instruction) = get(REDUCTION_CLASS, instr, NaN)
+function reduction_to_single_vector(x::Float64)
+    x == 1.0 ? :evadd : x == 2.0 ? :evmul : x == 3.0 ? :vand : x == 4.0 ? :vor : throw("Reduction not found.")
+end
+reduction_to_single_vector(x) = reduction_to_single_vector(reduction_instruction_class(x))
+function reduction_to_scalar(x::Float64)
+    x == 1.0 ? :vsum : x == 2.0 ? :vprod : x == 3.0 ? :vany : x == 4.0 ? :vall : 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 : throw("Reduction not found.")
+end
+reduction_scalar_combine(x) = reduction_scalar_combine(reduction_instruction_class(x))
+function reduction_combine_to(x::Float64)
+    x == 1.0 ? :reduce_to_add : x == 2.0 ? :reduce_to_prod : x == 3.0 ? :reduce_to_any : x == 4.0 ? :reduce_to_all : throw("Reduction not found.")
+end
+reduction_combine_to(x) = reduction_combine_to(reduction_instruction_class(x))
+function reduction_zero(x::Float64) 
+    x == 1.0 ? :zero : x == 2.0 ? :one : x == 3.0 ? :false : x == 4.0 ? :true : throw("Reduction not found.")
+end
+reduction_zero(x) = reduction_zero(reduction_instruction_class(x))
 
 const FUNCTIONSYMBOLS = Dict{Type{<:Function},Instruction}(
     typeof(+) => :(+),

src/lowering.jl

@@ -208,15 +208,14 @@ function reduce_range!(q::Expr, ls::LoopSet, Ulow::Int, Uhigh::Int)
         op = ls.operations[or]
         var = mangledvar(op)
         temp = gensym(var)
-        instr = op.instruction
-        instr = get(REDUCTION_TRANSLATION, instr, instr)
+        instr = Instruction(reduction_to_single_vector(op.instruction))
         reduce_range!(q, var, instr, Ulow, Uhigh)
     end
 end
 
 function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, U::Int)
     U == 1 && return nothing
-    instr = get(REDUCTION_TRANSLATION, instr, instr)
+    instr = Instruction(reduction_to_single_vector(instr))
     Uh2 = U
     iter = 0
     while true # combine vectors
@@ -226,8 +225,6 @@ function reduce_expr!(q::Expr, toreduct::Symbol, instr::Instruction, U::Int)
         Uh2 = Uh
         iter += 1; iter > 4 && throw("Oops! This seems to be excessive unrolling.")
     end
-    # reduce last vector
-    # push!(q.args, Expr(:(=), assignto, Expr(:call, reductfunc, Symbol(toreduct,:_0))))
     nothing
 end
 
@@ -403,7 +400,9 @@ function lower_compute!(
             # if op is an inner reduction, one of its parents will be the initialization of op
             # They will share the same `variable` field. The initialization may not have
             # unrolled in its loop dependencies, but (if opunrolled) op itself is, so we return true
-            parentsunrolled[p] = var === opp.variable ? true : (unrolled ∈ loopdependencies(opp))
+            # pu = (var === opp.variable || search_tree(parents(opp), var)) ? opunrolled : (unrolled ∈ loopdependencies(opp))
+            pu = unrolled ∈ loopdependencies(opp) || unrolled ∈ reducedchildren(opp)
+            parentsunrolled[p] = pu
         end
     else # maybe skip allocating this?
         parentsunrolled = fill(false, nparents)
@@ -662,8 +661,7 @@ end
 function initialize_outer_reductions!(
     q::Expr, op::Operation, Umin::Int, Umax::Int, W::Symbol, typeT::Symbol, vectorized::Symbol, suffix::Union{Symbol,Nothing} = nothing
 )
-    # T = op.elementbytes == 8 ? :Float64 : :Float32
-    z = Expr(:call, REDUCTION_ZERO[op.instruction], typeT)
+    z = Expr(:call, reduction_zero(op.instruction), typeT)
     if vectorized ∈ reduceddependencies(op)
         z = Expr(:call, lv(:vbroadcast), W, z)
     end
@@ -705,7 +703,7 @@ function reduce_expr!(q::Expr, ls::LoopSet, U::Int)
         mvar = mangledvar(op)
         instr = instruction(op)
         reduce_expr!(q, mvar, instr, U)
-        length(ls.opdict) == 0 || push!(q.args, Expr(:(=), var, Expr(:call, REDUCTION_SCALAR_COMBINE[instr], var, Symbol(mvar, 0))))
+        length(ls.opdict) == 0 || push!(q.args, Expr(:(=), var, Expr(:call, lv(reduction_scalar_combine(instr)), var, Symbol(mvar, 0))))
     end
 end
 function gc_preserve(ls::LoopSet, q::Expr)

src/operations.jl

@@ -94,6 +94,7 @@ struct Operation
     parents::Vector{Operation}
     ref::ArrayReferenceMeta
     mangledvariable::Symbol
+    reduced_children::Vector{Symbol}
     function Operation(
         identifier::Int,
         variable,
@@ -103,15 +104,16 @@ struct Operation
         dependencies = Symbol[],
         reduced_deps = Symbol[],
         parents = Operation[],
-        ref::ArrayReferenceMeta = NOTAREFERENCE
+        ref::ArrayReferenceMeta = NOTAREFERENCE,
+        reduced_children = Symbol[]
     )
         new(
             identifier, variable, elementbytes, instruction, node_type,
             convert(Vector{Symbol},dependencies),
             convert(Vector{Symbol},reduced_deps),
             convert(Vector{Operation},parents),
-            ref,
-            Symbol("##", variable, :_)
+            ref, Symbol("##", variable, :_),
+            reduced_children
         )
     end
 end
@@ -168,6 +170,7 @@ parents(op::Operation) = op.parents
 # children(op::Operation) = op.children
 loopdependencies(op::Operation) = op.dependencies
 reduceddependencies(op::Operation) = op.reduced_deps
+reducedchildren(op::Operation) = op.reduced_children
 identifier(op::Operation) = op.identifier + 1
 vptr(x::Symbol) = Symbol("##vptr##_", x)
 vptr(x::ArrayReference) = vptr(x.array)