search

TensorBFS / OMEinsumContractionOrders.jl

Tensor network contraction order optimizers for OMEinsum
MIT License
25 stars 3 forks source link

KaHyParBipartite or SABipartite as the initializer of TreeSA #35

Closed xptree closed 1 year ago

xptree commented 1 year ago

I noticed that there is an initializer argument in TreeSA algorithm. I wonder if it is possible to use the optcode returned by KaHyParBipartite or SABipartite as the initializer of TreeSA?

GiggleLiu commented 1 year ago

It is not supported yet. Seems to be reasonable to make the interface more generic. Will make the change in the next few days.

xptree commented 1 year ago

It is not supported yet. Seems to be reasonable to make the interface more generic. Will make the change in the next few days.

Thank you for your reply.

In the document of treesa, it says that initializer could be :greedy or :random

https://github.com/TensorBFS/OMEinsumContractionOrders.jl/blob/27a358c5e47cfde7c72a2765b7c4249b442672b6/src/treesa.jl#L157

However, in its internal implementation, I find a third choice of method == :specified

https://github.com/TensorBFS/OMEinsumContractionOrders.jl/blob/27a358c5e47cfde7c72a2765b7c4249b442672b6/src/treesa.jl#L236-L247

I wonder if the :specified option is a shortcut to enable a KaHyParBipartite contractor as the initializer? If the answer is yes, what is the proper way to make use of it.

GiggleLiu commented 1 year ago

Oh yes, you are right. The following should be the correct way of using it.

julia> using OMEinsum, KaHyPar

julia> code = ein"ijk,jkl,klm,lkm->il"
ijk, jkl, klm, lkm -> il

julia> tree = optimize_code(code, uniformsize(code, 2), KaHyParBipartite(sc_target=3))
ilk, kl -> il
├─ ijk, jkl -> ilk
│  ├─ ijk
│  └─ jkl
└─ klm, lkm -> kl
   ├─ klm
   └─ lkm

julia> optimize_code(tree, uniformsize(code, 2), TreeSA(; initializer=:specified))
SlicedEinsum{Char, DynamicNestedEinsum{Char}}(Char[], kl, lik -> il
├─ klm, lkm -> kl
│  ├─ klm
│  └─ lkm
└─ jkl, ijk -> lik
   ├─ jkl
   └─ ijk
)

Please let me know if it does not work.

xptree commented 1 year ago

Thank you! Seems that it works.

More background about my question:

I am integrating OMEinsumContractionOrders.jl into TensorCircuit as a customized contractor (https://github.com/tencent-quantum-lab/tensorcircuit/pull/107). Although I thought the initial status of SA may be crucial to its performance, I can not see significant difference between the two methods (:greedy and :specific) when contracting a google random circuit (circuit_n53_m20_s0_e0_pABCDCDAB),

OMEinsum contractor with kahypar init
------ contraction cost summary ------
log10[FLOPs]:  19.432  log2[SIZE]:  53  log2[WRITE]:  61.214

OMEinsum contractor with default greedy init
------ contraction cost summary ------
log10[FLOPs]:  19.435  log2[SIZE]:  53  log2[WRITE]:  61.262
GiggleLiu commented 1 year ago

🎉 Glad to know it works. I would suggest you to lower the initial temperature by e.g. setting β = 10:0.1:40 to make the initial vlaue effect more significant.

xptree commented 1 year ago

🎉 Glad to know it works. I would suggest you to lower the initial temperature by e.g. setting β = 10:0.1:40 to make the initial vlaue effect more significant.

A quick follow-up about your suggestion on tuning the temperature. I set β = 10:0.01:40 in kahypar initialization, and achieves better log10[FLOPs] comparing to either (1) greedy init or (2) kahypar init with default β = 0.01:0.01:15.

cotengra contractor
------ contraction cost summary ------
log10[FLOPs]:  19.561  log2[SIZE]:  53  log2[WRITE]:  61.980
OMEinsum contractor, greedy init
------ contraction cost summary ------
log10[FLOPs]:  19.442  log2[SIZE]:  53  log2[WRITE]:  61.235
OMEinsum contractor with kahypar init, by default β = 0.01:0.01:15
------ contraction cost summary ------
log10[FLOPs]:  19.436  log2[SIZE]:  53  log2[WRITE]:  61.271
OMEinsum contractor with kahypar init, lower the initial temperature, β = 10:0.01:40
------ contraction cost summary ------
log10[FLOPs]:  19.421  log2[SIZE]:  53  log2[WRITE]:  61.106

Thank you so much for your help!

GiggleLiu commented 1 year ago

Thank you very much for the update!

xptree commented 1 year ago

@GiggleLiu Hi, here are some further updates about this issue.

I tested the above method --- TreeSA with KaHyPar initialization --- with more google random circuits. But I found that for very few circuits, the method raises errors, and I have the following observations:

Do you have any ideas about the possible reason that lead to this error. I can provide more detailed data, such as input einsum, if necessary.

Thank you again! Wish you a very happy, healthy and successful Year of the Rabbit! 🐰🐰🐰

My julia script can be found at https://github.com/tencent-quantum-lab/tensorcircuit/blob/f87b758879003820e392aa4104ec157503313e93/examples/omeinsum_julia/omeinsum.jl, and more detailed stacktrace can be found as follows:

[ Info: `OMEinsumContractionOrders` loads `KaHyPar` module successfully.
ERROR: LoadError: TaskFailedException
Stacktrace:
 [1] wait
   @ ./task.jl:345 [inlined]
 [2] threading_run(fun::OMEinsumContractionOrders.var"#187#threadsfor_fun#110"{OMEinsumContractionOrders.var"#187#threadsfor_fun#107#111"{Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, Int64, Float64, Float64, Symbol, OMEinsumContractionOrders.MinSpaceOut, Int64, Vector{Any}, OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}, Vector{OMEinsumContractionOrders.Slicer}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{OMEinsumContractionOrders.ExprTree}, Vector{Float64}, Dict{Int64, Int64}, UnitRange{Int64}}}, static::Bool)
   @ Base.Threads ./threadingconstructs.jl:38
 [3] macro expansion
   @ ./threadingconstructs.jl:89 [inlined]
 [4] optimize_tree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, size_dict::Dict{Int64, Int64}; nslices::Int64, sc_target::Int64, βs::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, ntrials::Int64, niters::Int64, sc_weight::Float64, rw_weight::Float64, initializer::Symbol, greedy_method::OMEinsumContractionOrders.MinSpaceOut, greedy_nrepeat::Int64, fixed_slices::Vector{Any})
   @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:208
 [5] _optimize_code
   @ ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/interfaces.jl:60 [inlined]
 [6] optimize_code(code::OMEinsumContractionOrders.NestedEinsum{Int64}, size_dict::Dict{Int64, Int64}, optimizer::OMEinsumContractionOrders.TreeSA{Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, OMEinsumContractionOrders.GreedyMethod{OMEinsumContractionOrders.MinSpaceOut}, Any}, simplifier::Nothing, permute::Bool)
   @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/interfaces.jl:36
 [7] optimize_code(code::OMEinsum.DynamicNestedEinsum{Int64}, size_dict::Dict{Int64, Int64}, optimizer::OMEinsumContractionOrders.TreeSA{Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, OMEinsumContractionOrders.GreedyMethod{OMEinsumContractionOrders.MinSpaceOut}, Any}, simplifier::Nothing, permute::Bool) (repeats 2 times)
   @ OMEinsum ~/.julia/packages/OMEinsum/vgB8s/src/contractionorder.jl:31
 [8] main()
   @ Main ~/dev/tensorcircuit/examples/omeinsum_julia/omeinsum.jl:83
 [9] top-level scope
   @ ~/dev/tensorcircuit/examples/omeinsum_julia/omeinsum.jl:87

    nested task error: AssertionError: length(code.args) == 2
    Stacktrace:
      [1] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:457
      [2] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
      [3] iterate
        @ ./generator.jl:47 [inlined]
      [4] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:787
      [5] map
        @ ./abstractarray.jl:2961 [inlined]
      [6] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
      [7] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
      [8] iterate
        @ ./generator.jl:47 [inlined]
      [9] collect_to!(dest::Vector{OMEinsumContractionOrders.ExprTree}, itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}}, offs::Int64, st::Tuple{Int64, Int64})
        @ Base ./array.jl:845
     [10] collect_to_with_first!
        @ ./array.jl:823 [inlined]
     [11] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:797
     [12] map
        @ ./abstractarray.jl:2961 [inlined]
     [13] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
     [14] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
     [15] iterate
        @ ./generator.jl:47 [inlined]
     [16] collect_to!(dest::Vector{OMEinsumContractionOrders.ExprTree}, itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}}, offs::Int64, st::Tuple{Int64, Int64})
        @ Base ./array.jl:845
     [17] collect_to_with_first!
        @ ./array.jl:823 [inlined]
     [18] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:797
     [19] map
        @ ./abstractarray.jl:2961 [inlined]
     [20] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
     [21] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
     [22] iterate
        @ ./generator.jl:47 [inlined]
     [23] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:787
     [24] map
        @ ./abstractarray.jl:2961 [inlined]
     [25] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
     [26] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
     [27] iterate
        @ ./generator.jl:47 [inlined]
     [28] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:787
     [29] map
        @ ./abstractarray.jl:2961 [inlined]
     [30] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
     [31] (::OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}})(::Tuple{Int64, OMEinsumContractionOrders.NestedEinsum{Int64}})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:462
     [32] iterate
        @ ./generator.jl:47 [inlined]
     [33] collect_to!(dest::Vector{OMEinsumContractionOrders.ExprTree}, itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}}, offs::Int64, st::Tuple{Int64, Int64})
        @ Base ./array.jl:845
     [34] collect_to_with_first!
        @ ./array.jl:823 [inlined]
     [35] collect(itr::Base.Generator{Base.Iterators.Enumerate{Vector{OMEinsumContractionOrders.NestedEinsum}}, OMEinsumContractionOrders.var"#134#135"{OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}}})
        @ Base ./array.jl:797
     [36] map
        @ ./abstractarray.jl:2961 [inlined]
     [37] _exprtree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, labels::Dict{Int64, Int64})
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:458
     [38] _initializetree(code::OMEinsumContractionOrders.NestedEinsum{Int64}, size_dict::Dict{Int64, Int64}, method::Symbol; greedy_method::OMEinsumContractionOrders.MinSpaceOut, greedy_nrepeat::Int64)
        @ OMEinsumContractionOrders ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:244
     [39] macro expansion
        @ ~/.julia/packages/OMEinsumContractionOrders/dbGto/src/treesa.jl:210 [inlined]
     [40] (::OMEinsumContractionOrders.var"#187#threadsfor_fun#110"{OMEinsumContractionOrders.var"#187#threadsfor_fun#107#111"{Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, Int64, Float64, Float64, Symbol, OMEinsumContractionOrders.MinSpaceOut, Int64, Vector{Any}, OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}, Vector{OMEinsumContractionOrders.Slicer}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{OMEinsumContractionOrders.ExprTree}, Vector{Float64}, Dict{Int64, Int64}, UnitRange{Int64}}})(tid::Int64; onethread::Bool)
        @ OMEinsumContractionOrders ./threadingconstructs.jl:84
     [41] #187#threadsfor_fun
        @ ./threadingconstructs.jl:51 [inlined]
     [42] (::Base.Threads.var"#1#2"{OMEinsumContractionOrders.var"#187#threadsfor_fun#110"{OMEinsumContractionOrders.var"#187#threadsfor_fun#107#111"{Int64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, Int64, Float64, Float64, Symbol, OMEinsumContractionOrders.MinSpaceOut, Int64, Vector{Any}, OMEinsumContractionOrders.NestedEinsum{Int64}, Dict{Int64, Int64}, Vector{OMEinsumContractionOrders.Slicer}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{OMEinsumContractionOrders.ExprTree}, Vector{Float64}, Dict{Int64, Int64}, UnitRange{Int64}}}, Int64})()
        @ Base.Threads ./threadingconstructs.jl:30
in expression starting at ~/dev/tensorcircuit/examples/omeinsum_julia/omeinsum.jl:87