Search for parent refs as well as parent ops, also improve muladd-finder.

Author: chriselrod

src/add_compute.jl

@@ -86,6 +86,52 @@ function search_tree(opv::Vector{Operation}, var::Symbol) # relies on cycles bei
     end
     false
 end
+
+function update_for_ref_reduction!()
+    if varname(mpref) === var
+        id = findfirst(r -> r == mpref.mref, ls.refs_aliasing_syms)
+        mpref.varname = var = isnothing(id) ? var : ls.syms_aliasing_refs[id]
+        reduction_ind = ind
+        mergesetv!(deps, loopdependencies(add_load!(ls, argref, elementbytes)))
+    else
+        pushparent!(vparents, deps, reduceddeps, add_load!(ls, argref, elementbytes))
+    end
+end
+search_tree_for_ref(ls::LoopSet, opv::Vector{Operation}, ::Nothing, var::Symbol) = var, false
+function search_tree_for_ref(ls::LoopSet, opv::Vector{Operation}, mpref::ArrayReferenceMetaPosition, var::Symbol) # relies on cycles being forbidden
+    # isref, argref = tryrefconvert(ls, arg, elementbytes, varname(mpref))
+    # if isref
+    #     if mpref == argref
+    #         if varname(mpref) === var
+    #             id = findfirst(r -> r == mpref.mref, ls.refs_aliasing_syms)
+    #             mpref.varname = var = isnothing(id) ? var : ls.syms_aliasing_refs[id]
+    #             reduction_ind = ind
+    #             mergesetv!(deps, loopdependencies(add_load!(ls, argref, elementbytes)))
+    #         else
+    #             pushparent!(vparents, deps, reduceddeps, add_load!(ls, argref, elementbytes))
+    #         end
+    #     else
+    #         argref.varname = gensym!(ls, "tempload")
+    #         pushparent!(vparents, deps, reduceddeps, add_load!(ls, argref, elementbytes))
+    #     end
+    # else
+    #     add_parent!(vparents, deps, reduceddeps, ls, arg, elementbytes, position)
+    # end
+    for opp ∈ opv
+        if opp.ref == mpref.mref
+            if varname(mpref) === var
+                id = findfirst(r -> r == mpref.mref, ls.refs_aliasing_syms)
+                # @show var = isnothing(id) ? var : ls.syms_aliasing_refs[id]
+                mpref.varname = var = isnothing(id) ? var : ls.syms_aliasing_refs[id]
+                # @show mpref.varname
+                return var, true
+            end
+        end
+        var, found = search_tree_for_ref(ls, parents(opp), mpref, var)
+        found && return (var, found)
+    end
+    var, false
+end
 function search_tree(opv::Vector{Operation}, var::Operation) # relies on cycles being forbidden
     for opp ∈ opv
         opp === var && return true
@@ -281,26 +327,35 @@ function add_compute!(
         mergesetv!(newreduceddeps, reduceddeps)
         deps = newloopdeps; reduceddeps = newreduceddeps
     end
+    # @show reduction, search_tree(vparents, var) ex var vparents mpref get(ls.opdict, var, nothing) search_tree_for_ref(ls, vparents, mpref, var) # relies on cycles being forbidden
     op = if reduction || search_tree(vparents, var)
-        parent = ls.opdict[var]
-        setdiffv!(reduceddeps, deps, loopdependencies(parent))
-        # parent = getop(ls, var, elementbytes)
-        # if length(reduceddeps) == 0
-        if all(!in(deps), reduceddeps)
-            insert!(vparents, reduction_ind, parent)
-            mergesetv!(deps, loopdependencies(parent))
+        add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
+    else
+        var, found = search_tree_for_ref(ls, vparents, mpref, var)
+        if found
+            add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
+        else
             op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, vparents)
             pushop!(ls, op, var)
-        else
-            add_reduction_update_parent!(vparents, deps, reduceddeps, ls, parent, instr, reduction_ind, elementbytes)
         end
-    else
-        op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, vparents)
-        pushop!(ls, op, var)
     end
     # maybe_const_compute!(ls, op, elementbytes, position)
     op
 end
+function add_reduction!(ls::LoopSet, var::Symbol, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
+    parent = ls.opdict[var]
+    setdiffv!(reduceddeps, deps, loopdependencies(parent))
+    # parent = getop(ls, var, elementbytes)
+    # 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)
+        pushop!(ls, op, var)
+    else
+        add_reduction_update_parent!(vparents, deps, reduceddeps, ls, parent, instr, reduction_ind, elementbytes)
+    end    
+end
 
 function add_compute!(
     ls::LoopSet, LHS::Symbol, instr, vparents::Vector{Operation}, elementbytes::Int

src/vectorizationbase_compat/contract_pass.jl

@@ -5,6 +5,13 @@ function mulexprcost(ex::Expr)
     base = ex.head === :call ? 10 : 1
     base + length(ex.args)
 end
+function mul_fast_expr(args)
+    b = Expr(:call, :mul_fast)
+    for i ∈ 2:length(args)
+        push!(b.args, args[i])
+    end
+    b
+end
 function mulexpr(mulexargs)
     a = (mulexargs[1])::Union{Symbol,Expr,Number}
     if length(mulexargs) == 2
@@ -25,17 +32,17 @@ function mulexpr(mulexargs)
             return (c, Expr(:call, :mul_fast, a, b))
         end
     else
-        return (a, Expr(:call, :mul_fast, @view(mulexargs[2:end])...)::Expr)
+        return (a, mul_fast_expr(mulexargs))
     end
     a = (mulexargs[1])::Union{Symbol,Expr,Number}
     b = if length(mulexargs) == 2 # two arg mul
         (mulexargs[2])::Union{Symbol,Expr,Number}
     else
-        Expr(:call, :mul_fast, @view(mulexargs[2:end])...)::Expr
+        mul_fast_expr(mulexargs)
     end
     a, b
 end
-function append_args_skip!(call, args, i)
+function append_args_skip!(call, args, i, mod)
     for j ∈ eachindex(args)
         j == i && continue
         push!(call.args, args[j])
@@ -44,18 +51,29 @@ function append_args_skip!(call, args, i)
 end
 
 fastfunc(f) = get(VectorizationBase.FASTDICT, f, f)
-function make_fast!(call::Expr)
-    call.args[1] = fastfunc(first(call.args))
-    nothing
+function muladd_arguments!(argv, mod, f = first(argv))
+    if f === :*
+        argv[1] = :mul_fast
+    else
+        argv[1] = fastfunc(f)
+    end
+    for i ∈ 2:length(argv)
+        a = argv[i]
+        a isa Expr || continue
+        argv[i] = capture_muladd(a::Expr, mod)
+    end
 end
 
-function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool = false)
-    length(argv) < 3 && (make_fast!(call); return length(call.args) == 4, cnmul, csub)
+function recursive_muladd_search!(call, argv, mod, cnmul::Bool = false, csub::Bool = false)
+    if length(argv) < 3
+        muladd_arguments!(argv, mod)
+        return length(call.args) == 4, cnmul, csub
+    end
     fun = first(argv)
     isadd = fun === :+ || fun === :add_fast || fun === :vadd || (fun == :(Base.FastMath.add_fast))::Bool
     issub = fun === :- || fun === :sub_fast || fun === :vsub || (fun == :(Base.FastMath.sub_fast))::Bool
     if !(isadd | issub)
-        argv[1] = fastfunc(fun)
+        muladd_arguments!(argv, mod, fun)
         return length(call.args) == 4, cnmul, csub
     end
     exargs = @view(argv[2:end])
@@ -72,29 +90,29 @@ function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool =
                 if length(exargs) == 2
                     push!(call.args, exargs[3 -  i])
                 else
-                    push!(call.args, append_args_skip!(Expr(:call, :add_fast), exargs, i))
+                    push!(call.args, append_args_skip!(Expr(:call, :add_fast), exargs, i, mod))
                 end
                 if issub
                     csub = i == 1
                     cnmul = !csub
                 end
                 return true, cnmul, csub
             elseif isadd
-                found, cnmul, csub = recursive_muladd_search!(call, exa)
+                found, cnmul, csub = recursive_muladd_search!(call, exa, mod)
                 if found
                     if csub
                         call.args[4] = if length(exargs) == 2
                             Expr(:call, :sub_fast, exargs[3 - i], call.args[4])
                         else
-                            Expr(:call, :sub_fast, append_args_skip!(Expr(:call, :add_fast), exargs, i), call.args[4])
+                            Expr(:call, :sub_fast, append_args_skip!(Expr(:call, :add_fast), exargs, i, mod), call.args[4])
                         end
                     else
-                        call.args[4] = append_args_skip!(Expr(:call, :add_fast, call.args[4]), exargs, i)
+                        call.args[4] = append_args_skip!(Expr(:call, :add_fast, call.args[4]), exargs, i, mod)
                     end
                     return true, cnmul, false
                 end
             elseif issub
-                found, cnmul, csub = recursive_muladd_search!(call, exa)
+                found, cnmul, csub = recursive_muladd_search!(call, exa, mod)
                 if found
                     if i == 1
                         if csub
@@ -119,10 +137,11 @@ function recursive_muladd_search!(call, argv, cnmul::Bool = false, csub::Bool =
     length(call.args) == 4, cnmul, csub
 end
 
-function capture_muladd(ex::Expr, mod)
+function capture_a_muladd(ex::Expr, mod)
     call = Expr(:call, Symbol(""), Symbol(""), Symbol(""))
-    found, nmul, sub = recursive_muladd_search!(call, ex.args)
-    found || return ex
+    found, nmul, sub = recursive_muladd_search!(call, ex.args, mod)
+    found || return false, ex
+    # found || return ex
     # a, b, c = call.args[2], call.args[3], call.args[4]
     # call.args[2], call.args[3], call.args[4] = c, a, b
     f = if nmul && sub
@@ -139,7 +158,13 @@ function capture_muladd(ex::Expr, mod)
     else
         call.args[1] = Expr(:(.), mod, QuoteNote(f))#_fast))
     end
-    call
+    true, call
+end
+function capture_muladd(ex::Expr, mod)
+    while true
+        found, ex = capture_a_muladd(ex, mod)
+        found || return ex
+    end
 end
 
 contract_pass!(::Any, ::Any) = nothing

test/gemm.jl

@@ -572,7 +572,26 @@
     #         C[m,n] = Cmn_hi
     #     end
     # end
-    
+    function doublegemm!(du, u, mat)
+        @assert size(u, 1) == size(u, 2) == size(mat, 1) == size(mat, 2)
+        for i2 in 1:size(u, 2), i1 in 1:size(u, 1)
+            for sum_idx in 1:size(u, 1)
+                du[i1, i2] += mat[i1, sum_idx] * u[sum_idx, i2] + mat[i2, sum_idx] * u[i1, sum_idx]
+            end
+        end
+        return nothing
+    end
+
+    function doublegemmavx!(du, u, mat)
+        @assert size(u, 1) == size(u, 2) == size(mat, 1) == size(mat, 2)
+        @avx for i2 in 1:size(u, 2), i1 in 1:size(u, 1)
+            for sum_idx in 1:size(u, 1)
+                du[i1, i2] += mat[i1, sum_idx] * u[sum_idx, i2] + mat[i2, sum_idx] * u[i1, sum_idx]
+            end
+        end
+        return nothing
+    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
@@ -647,15 +666,22 @@
 # end
 
     for T ∈ (Float32, Float64, Int32, Int64)
+        TC = sizeof(T) == 4 ? Float32 : Float64
+        R = T <: Integer ? (T(-1000):T(1000)) : T
+        for M ∈ 48:54
+            C0 = zeros(TC, M, M); C1 = zeros(TC, M, M);
+            A = rand(R, M, M); B = rand(R, M, M);
+            doublegemm!(C0, A, B)
+            doublegemmavx!(C1, A, B)
+            @test C0 ≈ C1
+        end
     # let T = Int32
         # exceeds_time_limit() && break
         @show T, @__LINE__
         # M, K, N = 128, 128, 128;
         N = 69;
         for M ∈ 72:80, K ∈ 72:80
         # M, K, N = 73, 75, 69;
-            TC = sizeof(T) == 4 ? Float32 : Float64
-            R = T <: Integer ? (T(-1000):T(1000)) : T
             C = Matrix{TC}(undef, M, N);
             A = rand(R, M, K); B = rand(R, K, N);
             At = copy(A');