A new method for the greedy optimizer

Project.toml

@@ -5,9 +5,9 @@ version = "0.8.3"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
-BetterExp = "7cffe744-45fd-4178-b173-cf893948b8b7"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 
 [weakdeps]
@@ -18,7 +18,6 @@ KaHyParExt = ["KaHyPar"]
 
 [compat]
 AbstractTrees = "0.3, 0.4"
-BetterExp = "0.1"
 JSON = "0.21"
 KaHyPar = "0.3"
 Suppressor = "0.2"

src/OMEinsumContractionOrders.jl

@@ -2,14 +2,14 @@ module OMEinsumContractionOrders
 
 using JSON
 using SparseArrays
+using StatsBase
 using Base: RefValue
-using BetterExp
 using Base.Threads
 using AbstractTrees
 
 export CodeOptimizer, CodeSimplifier,
     KaHyParBipartite, GreedyMethod, TreeSA, SABipartite,
-    MinSpaceDiff, MinSpaceOut,
+    MinSpaceDiff, MinSpaceOut, HyperGreedy,
     MergeGreedy, MergeVectors,
     uniformsize,
     simplify_code, optimize_code, optimize_permute,

src/greedy.jl

@@ -5,7 +5,20 @@ end
 
 struct MinSpaceOut end
 struct MinSpaceDiff end
+struct HyperGreedy{T}
+    t::T
+    HyperGreedy(t::T = 1.0) where T = new{T}(t)
+end
+
+
+"""
+    LegInfo{ET}
 
+    A struct to store the information of legs in a pairwise contraction between vertices `vi` and `vj`.
+    *`l*` are the legs that are not connected to the other vertex (internal legs)
+    *`l0*` are the legs that are connected to other vertices (external legs)
+    *`*1` are the legs that are connected only to `vi`, and `*2` are the legs that are connected only to `vj`, and `*12` are the legs that are connected to both.
+"""
 struct LegInfo{ET}
     l1::Vector{ET}
     l2::Vector{ET}
@@ -50,15 +63,19 @@ ae, ak -> ea
 """
 function tree_greedy(incidence_list::IncidenceList{VT,ET}, log2_edge_sizes; method=MinSpaceOut(), nrepeat=10) where {VT,ET}
     @assert nrepeat >= 1
-    best_tree, best_tcs, best_scs = _tree_greedy(incidence_list, log2_edge_sizes; method=method)
-    best_tc, best_sc = log2sumexp2(best_tcs), maximum(best_scs)
-    for _ = 1:nrepeat-1
-        tree, tcs, scs = _tree_greedy(incidence_list, log2_edge_sizes; method=method)
-        tc, sc = log2sumexp2(tcs), maximum(scs)
-        if sc < best_sc || (sc <= best_sc && tc < best_tc)
-            best_tcs, best_scs, best_tc, best_sc, best_tree = tcs, scs, tc, sc, tree
-        end
+
+    results = Vector{Tuple{ContractionTree, Vector{Float64}, Vector{Float64}}}(undef, nrepeat)
+
+    @threads for i = 1:nrepeat
+        results[i] = _tree_greedy(incidence_list, log2_edge_sizes; method=method)
     end
+
+    best_sc = minimum([maximum(r[3]) for r in results])
+    possible_ids = findall(x -> maximum(x[3]) == best_sc, results)
+    possible_results = results[possible_ids]
+
+    best_tree, best_tcs, best_scs = results[argmin([log2sumexp2(r[2]) for r in possible_results])]
+
     return best_tree, best_tcs, best_scs
 end
 
@@ -80,7 +97,11 @@ function _tree_greedy(incidence_list::IncidenceList{VT,ET}, log2_edge_sizes; met
             vpool = collect(vertices(incidence_list))
             pair = minmax(vpool[1], vpool[2])  # to prevent empty intersect
         else
-            pair = find_best_cost(cost_values)
+            if method isa HyperGreedy
+                pair = sample_best_cost(cost_values, method.t)
+            else
+                pair = find_best_cost(cost_values)
+            end
         end
         log2_tc_step, sc, code = contract_pair!(incidence_list, pair..., log2_edge_sizes)
         push!(log2_tcs, log2_tc_step)
@@ -153,6 +174,15 @@ function find_best_cost(cost_values::Dict{PT}) where PT
     return rand(pairs)
 end
 
+function sample_best_cost(cost_values::Dict{PT}, t::T) where {PT, T}
+    length(cost_values) < 1 && error("cost value information missing")
+    vals = [v for v in values(cost_values)]
+    prob = exp.( - vals ./ t)
+    prob = prob ./ sum(prob)
+    vc = [k for (k, v) in cost_values]
+    sample(vc, Weights(prob))
+end
+
 function analyze_contraction(incidence_list::IncidenceList{VT,ET}, vi::VT, vj::VT) where {VT,ET}
     ei = edges(incidence_list, vi)
     ej = edges(incidence_list, vj)
@@ -185,7 +215,7 @@ function analyze_contraction(incidence_list::IncidenceList{VT,ET}, vi::VT, vj::V
     return LegInfo(leg1, leg2, leg12, leg01, leg02, leg012)
 end
 
-function greedy_loss(::MinSpaceOut, incidence_list, log2_edge_sizes, vi, vj)
+function greedy_loss(::Union{MinSpaceOut, HyperGreedy}, incidence_list, log2_edge_sizes, vi, vj)
     log2dim(legs) = isempty(legs) ? 0 : sum(l->log2_edge_sizes[l], legs)  # for 1.5, you need this patch because `init` kw is not allowed.
     legs = analyze_contraction(incidence_list, vi, vj)
     log2dim(legs.l01)+log2dim(legs.l02)+log2dim(legs.l012)
@@ -298,9 +328,10 @@ end
 
 The fast but poor greedy optimizer. Input arguments are
 
-* `method` is `MinSpaceDiff()` or `MinSpaceOut`.
+* `method` is `MinSpaceDiff()`, `MinSpaceOut()` or `HyperGreedy(α=0.0, t=1.0)`.
     * `MinSpaceOut` choose one of the contraction that produces a minimum output tensor size,
     * `MinSpaceDiff` choose one of the contraction that decrease the space most.
+    * `HyperGreedy` is a hyperparameterized greedy method, `t` is the temperature.
 * `nrepeat` is the number of repeatition, returns the best contraction order.
 """
 Base.@kwdef struct GreedyMethod{MT} <: CodeOptimizer

src/sa.jl

@@ -72,9 +72,9 @@ function bipartite_sc(bipartiter::SABipartite, adj::SparseMatrixCSC, vertices, l
             newloss = compute_loss(sc_ti, sc_tinew, state.group_sizes[ti]-1, state.group_sizes[3-ti]+1)
             sc_ti0, sc_tinew0 = state.group_scs[ti], state.group_scs[3-ti]
             accept = if max(sc_ti0, sc_tinew0) <= bipartiter.sc_target
-                max(sc_ti, sc_tinew) <= bipartiter.sc_target && (rand() < BetterExp.exp2(β*(state.loss[] - newloss)))
+                max(sc_ti, sc_tinew) <= bipartiter.sc_target && (rand() < exp2(β*(state.loss[] - newloss)))
             else
-                rand() < BetterExp.exp2(-β*(max(sc_ti, sc_tinew) - max(sc_ti0, sc_tinew0)))
+                rand() < exp2(-β*(max(sc_ti, sc_tinew) - max(sc_ti0, sc_tinew0)))
             end
             accept && update_state!(state, adjt, vertices, idxi, sc_ti, sc_tinew, newloss)
         end

test/greedy.jl

@@ -1,5 +1,6 @@
 using OMEinsumContractionOrders
 using OMEinsumContractionOrders: analyze_contraction, contract_pair!, evaluate_costs, contract_tree!, log2sumexp2, parse_tree
+using OMEinsumContractionOrders: IncidenceList, analyze_contraction, LegInfo, tree_greedy, parse_eincode, optimize_greedy
 using TropicalNumbers
 
 using Test, Random
@@ -34,27 +35,28 @@ end
     contract_tree!(copy(incidence_list), tree, log2_edge_sizes, tcs_, scs_)
     @test all((log2sumexp2(tcs_), maximum(scs_)) .<= (log2(exp2(10)+exp2(16)+exp2(15)+exp2(9)), 11))
     vertices = ['A', 'B', 'C', 'D', 'E']
-    optcode1 = parse_eincode(StaticEinCode, incidence_list, tree, vertices=vertices)
+    optcode1 = parse_eincode(incidence_list, tree, vertices=vertices)
     @test optcode1 isa OMEinsumContractionOrders.NestedEinsum
     tree2 = parse_tree(optcode1, vertices)
     @test tree2 == tree
 
-    eincode = ein"ab,acd,bcef,e,df->"
+    eincode = OMEinsumContractionOrders.EinCode([['a', 'b'], ['a', 'c', 'd'], ['b', 'c', 'e', 'f'], ['e'], ['d', 'f']], Vector{Char}())
     size_dict = Dict([c=>(1<<i) for (i,c) in enumerate(['a', 'b', 'c', 'd', 'e', 'f'])]...)
     Random.seed!(2)
     optcode2 = optimize_greedy(eincode, size_dict) 
-    tc, sc = timespace_complexity(optcode2, edge_sizes)
+    cc = contraction_complexity(optcode2, edge_sizes)
     # test flop
-    @test tc ≈ log2(flop(optcode2, edge_sizes))
-    @test flop(ein"i->", Dict('i'=>4)) == 4
-    @test 16 <= tc <= log2(exp2(10)+exp2(16)+exp2(15)+exp2(9))
-    @test sc == 11
+    @test cc.tc ≈ log2(flop(optcode2, edge_sizes))
+    @test flop(OMEinsumContractionOrders.EinCode([['i']], Vector{Char}()), Dict('i'=>4)) == 4
+    @test 16 <= cc.tc <= log2(exp2(10)+exp2(16)+exp2(15)+exp2(9))
+    @test cc.sc == 11
     @test optcode1 == optcode2
-    eincode3 = ein"(ab,acd),bcef,e,df->"
-    Random.seed!(2)
-    optcode3 = optimize_greedy(eincode3, size_dict) 
-    tc, sc = timespace_complexity(optcode3, edge_sizes)
-    @test 16 <= tc <= log2(exp2(10)+exp2(16)+exp2(15)+exp2(9)+1e-8)
+
+    # eincode3 = ein"(ab,acd),bcef,e,df->"
+    # Random.seed!(2)
+    # optcode3 = optimize_greedy(eincode3, size_dict) 
+    # tc, sc = timespace_complexity(optcode3, edge_sizes)
+    # @test 16 <= tc <= log2(exp2(10)+exp2(16)+exp2(15)+exp2(9)+1e-8)
 end
 
 @testset "fullerene" begin
@@ -73,61 +75,45 @@ end
         return res
     end
 
+    # flatten nested einsum
+    function _flatten(code::OMEinsumContractionOrders.NestedEinsum, iy=nothing)
+        isleaf(code) && return [tensorindex(code)=>iy]
+        sibs = siblings(code)
+        ixs = []
+        for i=1:length(sibs)
+            append!(ixs, _flatten(sibs[i], (rootcode(code).ixs)[i]))
+        end
+        return ixs
+    end
+
+    flatten(code::OMEinsumContractionOrders.EinCode) = code
+    function flatten(code::OMEinsumContractionOrders.NestedEinsum{LT}) where LT
+        ixd = Dict(_flatten(code))
+        OMEinsumContractionOrders.EinCode([ixd[i] for i=1:length(ixd)], collect((code.eins).iy))
+    end
+
+    isleaf(ne::OMEinsumContractionOrders.NestedEinsum) = ne.tensorindex != -1
+    siblings(ne::OMEinsumContractionOrders.NestedEinsum) = ne.args
+    tensorindex(ne::OMEinsumContractionOrders.NestedEinsum) = ne.tensorindex
+    rootcode(ne::OMEinsumContractionOrders.NestedEinsum) = ne.eins
+
     c60_xy = fullerene()
-    c60_edges = [(i,j) for (i,(i2,j2,k2)) in enumerate(c60_xy), (j,(i1,j1,k1)) in enumerate(c60_xy) if i<j && (i2-i1)^2+(j2-j1)^2+(k2-k1)^2 < 5.0]
-    code = EinCode((c60_edges..., [(i,) for i=1:60]...), ())
+    c60_edges = [[i,j] for (i,(i2,j2,k2)) in enumerate(c60_xy), (j,(i1,j1,k1)) in enumerate(c60_xy) if i<j && (i2-i1)^2+(j2-j1)^2+(k2-k1)^2 < 5.0]
+    code = OMEinsumContractionOrders.EinCode(vcat(c60_edges, [[i] for i=1:60]), Vector{Int}())
     size_dict = Dict([i=>2 for i in 1:60])
     log2_edge_sizes = Dict([i=>1 for i in 1:60])
     edge_sizes = Dict([i=>2 for i in 1:60])
-    tc, sc = timespace_complexity(code, edge_sizes)
-    @test tc == 60
-    @test sc == 0
+    cc = contraction_complexity(code, edge_sizes)
+    @test cc.tc == 60
+    @test cc.sc == 0
     optcode = optimize_greedy(code, size_dict)
-    tc2, sc2 = timespace_complexity(optcode, edge_sizes)
-    @test sc2 == 10
-    xs = vcat([TropicalF64.([-1 1; 1 -1]) for i=1:90], [TropicalF64.([0, 0]) for i=1:60])
-    @test OMEinsum.flatten(optcode) == code
-    @test OMEinsum.flatten(code) == code
-    @test optcode(xs...)[].n == 66
-end
-
-@testset "regression test" begin
-    code = ein"i->"
-    optcode = optimize_greedy(code, Dict('i'=>3))
-    @test optcode isa NestedEinsum
-    x = randn(3)
-    @test optcode(x) ≈ code(x)
-
-    code = ein"i,j->"
-    optcode = optimize_greedy(code, Dict('i'=>3, 'j'=>3))
-    @test optcode isa NestedEinsum
-    x = randn(3)
-    y = randn(3)
-    @test optcode(x, y) ≈ code(x, y)
-
-    code = ein"ij,jk,kl->ijl"
-    optcode = optimize_greedy(code, Dict('i'=>3, 'j'=>3, 'k'=>3, 'l'=>3))
-    @test optcode isa NestedEinsum
-    a, b, c = [rand(3,3) for i=1:3]
-    @test optcode(a, b, c) ≈ code(a, b, c)
-end
-
-@testset "constructing contraction tree manually" begin
-    code = ein"ij,jk,kl->ijl"
-    dcode = DynamicEinCode(ein"ij,jk,kl->ijl")
-    a, b, c = randn(2, 2), randn(2,2), randn(2,2)
-    ne1 = parse_nested(code, ContractionTree(ContractionTree(1, 2), 3))
-    ne2 = parse_nested(dcode, ContractionTree(ContractionTree(1, 2), 3))
-    ne3 = parse_nested(code, ContractionTree(ContractionTree(1, 3), 2))
-    @test typeof(ne1) == NestedEinsum{StaticEinCode}
-    @test typeof(ne2) == NestedEinsum{DynamicEinCode{Char}}
-    @test ne1(a,b,c) ≈ ein"(ij,jk),kl->ijl"(a,b,c)
-    @test ne2(a,b,c) ≈ ein"(ij,jk),kl->ijl"(a,b,c)
-    @test ne3(a,b,c) ≈ ein"(ij,jk),kl->ijl"(a,b,c)
-    @test flop(code, Dict([l=>2 for l in uniquelabels(code)])) == 2^4
-    @test flop(ne1, Dict([l=>2 for l in uniquelabels(code)])) == 2^4 + 2^3
-    @test flop(ne2, Dict([l=>2 for l in uniquelabels(code)])) == 2^4 + 2^3
+    cc2 = contraction_complexity(optcode, edge_sizes)
+    @test cc2.sc == 10
+    @test flatten(optcode) == code
+    @test flatten(code) == code
 
-    # label elimination order
-    @test label_elimination_order(ein"(ij,jk),kl->il") == ['j', 'k']
+    optcode_hyper = optimize_greedy(code, size_dict, method = HyperGreedy(0.1))
+    cc3 = contraction_complexity(optcode_hyper, edge_sizes)
+    @test cc3.sc == 10
+    @test flatten(optcode_hyper) == code
 end

test/runtests.jl

@@ -1,6 +1,10 @@
 using OMEinsumContractionOrders
 using Test
 
+@testset "greedy" begin
+    include("greedy.jl")
+end
+
 @testset "kahypar" begin
     include("kahypar.jl")
 end