From cc2295faae60f813f17d6dfdbf539ced8d138942 Mon Sep 17 00:00:00 2001 From: Jinguo Liu Date: Sun, 24 Sep 2023 01:33:51 +0800 Subject: [PATCH] Jg/switch to extension (#39) * switch to extensions from Requires * require julia 1.9 * bump version * update README and improve complexity printing --- CITATION.bib | 51 +++++++++++ Project.toml | 16 ++-- README.md | 147 +++++++++---------------------- ext/KaHyParExt.jl | 30 +++++++ src/OMEinsumContractionOrders.jl | 9 -- src/complexity.jl | 6 +- src/kahypar.jl | 25 +----- src/treesa.jl | 2 +- test/interfaces.jl | 4 +- 9 files changed, 142 insertions(+), 148 deletions(-) create mode 100644 CITATION.bib create mode 100644 ext/KaHyParExt.jl diff --git a/CITATION.bib b/CITATION.bib new file mode 100644 index 0000000..c4f0027 --- /dev/null +++ b/CITATION.bib @@ -0,0 +1,51 @@ +@misc{Liu2022, + doi = {10.48550/ARXIV.2205.03718}, + url = {https://arxiv.org/abs/2205.03718}, + author = {Liu, Jin-Guo and Gao, Xun and Cain, Madelyn and Lukin, Mikhail D. and Wang, Sheng-Tao}, + keywords = {Statistical Mechanics (cond-mat.stat-mech), FOS: Physical sciences, FOS: Physical sciences}, + title = {Computing solution space properties of combinatorial optimization problems via generic tensor networks}, + publisher = {arXiv}, + year = {2022}, + copyright = {Creative Commons Attribution 4.0 International} +} +@misc{Pan2021, + title={Simulating the Sycamore quantum supremacy circuits}, + author={Feng Pan and Pan Zhang}, + year={2021}, + eprint={2103.03074}, + archivePrefix={arXiv}, + primaryClass={quant-ph} +} +@Article{10.21468/SciPostPhys.7.5.060, + title={{Fast counting with tensor networks}}, + author={Stefanos Kourtis and Claudio Chamon and Eduardo R. Mucciolo and Andrei E. Ruckenstein}, + journal={SciPost Phys.}, + volume={7}, + issue={5}, + pages={60}, + year={2019}, + publisher={SciPost}, + doi={10.21468/SciPostPhys.7.5.060}, + url={https://scipost.org/10.21468/SciPostPhys.7.5.060}, +} +@article{Gray2021, + title={Hyper-optimized tensor network contraction}, + volume={5}, + ISSN={2521-327X}, + url={http://dx.doi.org/10.22331/q-2021-03-15-410}, + DOI={10.22331/q-2021-03-15-410}, + journal={Quantum}, + publisher={Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften}, + author={Gray, Johnnie and Kourtis, Stefanos}, + year={2021}, + month={Mar}, + pages={410} +} +@misc{kalachev2021recursive, + title={Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits}, + author={Gleb Kalachev and Pavel Panteleev and Man-Hong Yung}, + year={2021}, + eprint={2108.05665}, + archivePrefix={arXiv}, + primaryClass={quant-ph} +} diff --git a/Project.toml b/Project.toml index c7cecd0..675ae4f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,31 +1,35 @@ name = "OMEinsumContractionOrders" uuid = "6f22d1fd-8eed-4bb7-9776-e7d684900715" authors = ["Jin-Guo Liu", "Pan Zhang"] -version = "0.8.2" +version = "0.8.3" [deps] AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" BetterExp = "7cffe744-45fd-4178-b173-cf893948b8b7" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -Requires = "ae029012-a4dd-5104-9daa-d747884805df" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" +[weakdeps] +KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880" + +[extensions] +KaHyParExt = ["KaHyPar"] + [compat] AbstractTrees = "0.3, 0.4" BetterExp = "0.1" JSON = "0.21" -Requires = "1" +KaHyPar = "0.3" Suppressor = "0.2" -julia = "1" +julia = "1.9" [extras] Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" -KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880" OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334" [targets] -test = ["Test", "Random", "KaHyPar", "Graphs", "TropicalNumbers", "OMEinsum"] +test = ["Test", "Random", "Graphs", "TropicalNumbers", "OMEinsum", "KaHyPar"] diff --git a/README.md b/README.md index 5687411..1b51d51 100644 --- a/README.md +++ b/README.md @@ -1,26 +1,22 @@ # OMEinsumContractionOrders -This package provides `optimize_code` function for finding tensor network contraction orders. + +OMEinsumContractionOrders is a Julia package that provides an `optimize_code` function for finding optimal contraction orders for tensor networks. It is designed to work with multiple tensor network packages, such as: [OMEinsum.jl](https://github.com/under-Peter/OMEinsum.jl/issues) package and [ITensorNetworks.jl](https://github.com/mtfishman/ITensorNetworks.jl). [![Build Status](https://github.com/TensorBFS/OMEinsumContractionOrders.jl/workflows/CI/badge.svg)](https://github.com/TensorBFS/OMEinsumContractionOrders.jl/actions) [![codecov](https://codecov.io/gh/TensorBFS/OMEinsumContractionOrders.jl/branch/master/graph/badge.svg?token=BwaF0cV6q1)](https://codecov.io/gh/TensorBFS/OMEinsumContractionOrders.jl) ## Installation -

-OMEinsumContractionOrders is a Julia Language package. To install OMEinsumContractionOrders, - please open - Julia's interactive session (known as REPL) and press ] key in the REPL to use the package mode, then type the following command -

-For stable release +To install OMEinsumContractionOrders, please follow these steps: -```julia -pkg> add OMEinsumContractionOrders -pkg> add KaHyPar -``` -The `KaHyPar` package (used in `KaHyParBipartite` optimizer) is optional because some one may have the binary issue, check [this](https://github.com/kahypar/KaHyPar.jl/issues/12) and [this](https://github.com/kahypar/KaHyPar.jl/issues/19). +1. Open Julia's interactive session (known as [REPL](https://docs.julialang.org/en/v1/manual/getting-started/)) by typing `julia` in your terminal. +2. Press the ] key in the REPL to enter the package mode. +3. Type `add OMEinsumContractionOrders` to install the stable release of the package. +4. (Optional) If you want to use the `KaHyParBipartite` optimizer, which is based on the KaHyPar library, type `add KaHyPar`. Note that this step is optional because some users may have issues with the binary dependencies of KaHyPar (please check issues: [this](https://github.com/kahypar/KaHyPar.jl/issues/12) and [this](https://github.com/kahypar/KaHyPar.jl/issues/19)). + +## Example 1: Use it directly +The contraction order optimizer is implemented in the `optimize_code` function. It takes three arguments: `code`, `size`, and `optimizer`. The `code` argument is the [einsum notation](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html) to be optimized. The `size` argument is the size of the variables in the einsum notation. The `optimizer` argument is the optimizer to be used. The `optimize_code` function returns the optimized contraction order. One can use `contraction_complexity` function to get the time, space and rewrite complexity of returned contraction order. -## Example -Optimize a contraction order ```julia julia> using OMEinsumContractionOrders, Graphs, KaHyPar @@ -44,23 +40,35 @@ julia> optcode_kahypar = optimize_code(code, uniformsize(code, 2), julia> optcode_sa = optimize_code(code, uniformsize(code, 2), SABipartite(sc_target=30, max_group_size=50)); -julia> timespace_complexity(code, uniformsize(code, 2)) -(200.0, 0.0) - -julia> timespace_complexity(optcode_kahypar, uniformsize(code, 2)) -(39.00774639886569, 28.0) - -julia> timespace_complexity(optcode_sa, uniformsize(code, 2)) -(39.09524927371961, 28.0) - -julia> timespace_complexity(optcode_tree, uniformsize(code, 2)) -(31.13883492534964, 27.0) - -julia> timespace_complexity(optcode_tree_with_slice, uniformsize(code, 2)) -(31.292828948918775, 22.0) +julia> contraction_complexity(code, uniformsize(code, 2)) +Time complexity: 2^200.0 +Space complexity: 2^0.0 +Read-write complexity: 2^10.644757592516257 + +julia> contraction_complexity(optcode_kahypar, uniformsize(code, 2)) +Time complexity: 2^39.5938886486877 +Space complexity: 2^28.0 +Read-write complexity: 2^30.39890775966298 + +julia> contraction_complexity(optcode_sa, uniformsize(code, 2)) +Time complexity: 2^41.17129641027078 +Space complexity: 2^29.0 +Read-write complexity: 2^31.493976190321106 + +julia> contraction_complexity(optcode_tree, uniformsize(code, 2)) +Time complexity: 2^35.06468305863757 +Space complexity: 2^28.0 +Read-write complexity: 2^30.351552349259258 + +julia> contraction_complexity(optcode_tree_with_slice, uniformsize(code, 2)) +Time complexity: 2^33.70760100663681 +Space complexity: 2^24.0 +Read-write complexity: 2^32.17575935629581 ``` -It is already a part of [`OMEinsum`](https://github.com/under-Peter/OMEinsum.jl): +## Example 2: Use it in `OMEinsum` + +`OMEinsumContractionOrders` is shipped with [`OMEinsum`](https://github.com/under-Peter/OMEinsum.jl) package. You can use it to optimize the contraction order of an `OMEinsum` expression. ```julia julia> using OMEinsum @@ -68,7 +76,7 @@ julia> using OMEinsum julia> code = ein"ij, jk, kl, il->" ij, jk, kl, il -> -julia> optimize_code(code, uniformsize(code, 2), TreeSA()) +julia> optimized_code = optimize_code(code, uniformsize(code, 2), TreeSA()) SlicedEinsum{Char, NestedEinsum{DynamicEinCode{Char}}}(Char[], ki, ki -> ├─ jk, ij -> ki │ ├─ jk @@ -79,86 +87,15 @@ SlicedEinsum{Char, NestedEinsum{DynamicEinCode{Char}}}(Char[], ki, ki -> ) ``` -## Acknowledge OMEinsum/OMEinsumContractionOrders -Since we do not have a paper to cite, we wish anyone who finds this package useful in his research can post a comment under this issue: -https://github.com/TensorBFS/OMEinsumContractionOrders.jl/issues/21 - ## References -If you find this package useful in your research, please cite the following papers -``` -@misc{Liu2022, - doi = {10.48550/ARXIV.2205.03718}, - url = {https://arxiv.org/abs/2205.03718}, - author = {Liu, Jin-Guo and Gao, Xun and Cain, Madelyn and Lukin, Mikhail D. and Wang, Sheng-Tao}, - keywords = {Statistical Mechanics (cond-mat.stat-mech), FOS: Physical sciences, FOS: Physical sciences}, - title = {Computing solution space properties of combinatorial optimization problems via generic tensor networks}, - publisher = {arXiv}, - year = {2022}, - copyright = {Creative Commons Attribution 4.0 International} -} -``` - -To credit the `KaHyParBipartite` and `SABipartite` method, -``` -@misc{Pan2021, - title={Simulating the Sycamore quantum supremacy circuits}, - author={Feng Pan and Pan Zhang}, - year={2021}, - eprint={2103.03074}, - archivePrefix={arXiv}, - primaryClass={quant-ph} -} -``` - -To credit the `KaHyParBipartite` method, -``` -@Article{10.21468/SciPostPhys.7.5.060, - title={{Fast counting with tensor networks}}, - author={Stefanos Kourtis and Claudio Chamon and Eduardo R. Mucciolo and Andrei E. Ruckenstein}, - journal={SciPost Phys.}, - volume={7}, - issue={5}, - pages={60}, - year={2019}, - publisher={SciPost}, - doi={10.21468/SciPostPhys.7.5.060}, - url={https://scipost.org/10.21468/SciPostPhys.7.5.060}, -} -``` - -``` -@article{Gray2021, - title={Hyper-optimized tensor network contraction}, - volume={5}, - ISSN={2521-327X}, - url={http://dx.doi.org/10.22331/q-2021-03-15-410}, - DOI={10.22331/q-2021-03-15-410}, - journal={Quantum}, - publisher={Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften}, - author={Gray, Johnnie and Kourtis, Stefanos}, - year={2021}, - month={Mar}, - pages={410} -} -``` - -To credit the `TreeSA` method, -``` -@misc{kalachev2021recursive, - title={Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits}, - author={Gleb Kalachev and Pavel Panteleev and Man-Hong Yung}, - year={2021}, - eprint={2108.05665}, - archivePrefix={arXiv}, - primaryClass={quant-ph} -} -``` +If you find this package useful in your research, please cite the *relevant* papers in [CITATION.bib](CITATION.bib). ## Multi-GPU computation -Check this Gist: +Please check this Gist: https://gist.github.com/GiggleLiu/d5b66c9883f0c5df41a440589983ab99 ## Authors -Jin-Guo Liu and Pan Zhang + +OMEinsumContractionOrders was developed by Jin-Guo Liu and Pan Zhang. \ No newline at end of file diff --git a/ext/KaHyParExt.jl b/ext/KaHyParExt.jl new file mode 100644 index 0000000..5040db6 --- /dev/null +++ b/ext/KaHyParExt.jl @@ -0,0 +1,30 @@ +module KaHyParExt + +using OMEinsumContractionOrders: KaHyParBipartite, SparseMatrixCSC, group_sc, induced_subhypergraph, convert2int +import KaHyPar +import OMEinsumContractionOrders: bipartite_sc +using Suppressor: @suppress + +function bipartite_sc(bipartiter::KaHyParBipartite, adj::SparseMatrixCSC, vertices, log2_sizes) + n_v = length(vertices) + subgraph, remaining_edges = induced_subhypergraph(adj, vertices) + hypergraph = KaHyPar.HyperGraph(subgraph, ones(n_v), convert2int(log2_sizes[remaining_edges])) + local parts + min_sc = 999999 + for imbalance in bipartiter.imbalances + parts = @suppress KaHyPar.partition(hypergraph, 2; imbalance=imbalance, configuration=:edge_cut) + part0 = vertices[parts .== 0] + part1 = vertices[parts .== 1] + sc0, sc1 = group_sc(adj, part0, log2_sizes), group_sc(adj, part1, log2_sizes) + sc = max(sc0, sc1) + min_sc = min(sc, min_sc) + @debug "imbalance $imbalance: sc = $sc, group = ($(length(part0)), $(length(part1)))" + if sc <= bipartiter.sc_target + return part0, part1 + end + end + error("fail to find a valid partition for `sc_target = $(bipartiter.sc_target)`, got minimum value `$min_sc` (imbalances = $(bipartiter.imbalances))") +end + +@info "`OMEinsumContractionOrders` loads `KaHyParExt` extension successfully." +end \ No newline at end of file diff --git a/src/OMEinsumContractionOrders.jl b/src/OMEinsumContractionOrders.jl index 7a9c3fd..2e2ee89 100644 --- a/src/OMEinsumContractionOrders.jl +++ b/src/OMEinsumContractionOrders.jl @@ -5,17 +5,8 @@ using SparseArrays using Base: RefValue using BetterExp using Base.Threads -using Suppressor: @suppress using AbstractTrees -using Requires -function __init__() - @require KaHyPar="2a6221f6-aa48-11e9-3542-2d9e0ef01880" begin - using .KaHyPar - @info "`OMEinsumContractionOrders` loads `KaHyPar` module successfully." - end -end - export CodeOptimizer, CodeSimplifier, KaHyParBipartite, GreedyMethod, TreeSA, SABipartite, MinSpaceDiff, MinSpaceOut, diff --git a/src/complexity.jl b/src/complexity.jl index 36f8043..8411f14 100644 --- a/src/complexity.jl +++ b/src/complexity.jl @@ -172,9 +172,9 @@ struct ContractionComplexity end function Base.show(io::IO, cc::ContractionComplexity) - print(io, "Time complexity (number of element-wise multiplications) = 2^$(cc.tc) -Space complexity (number of elements in the largest intermediate tensor) = 2^$(cc.sc) -Read-write complexity (number of element-wise read and write) = 2^$(cc.rwc)") + print(io, "Time complexity: 2^$(cc.tc) +Space complexity: 2^$(cc.sc) +Read-write complexity: 2^$(cc.rwc)") end Base.iterate(cc::ContractionComplexity) = Base.iterate((cc.tc, cc.sc, cc.rwc)) Base.iterate(cc::ContractionComplexity, state) = Base.iterate((cc.tc, cc.sc, cc.rwc), state) diff --git a/src/kahypar.jl b/src/kahypar.jl index 59e3e06..4062f22 100644 --- a/src/kahypar.jl +++ b/src/kahypar.jl @@ -35,28 +35,9 @@ function convert2int(sizes::AbstractVector) round.(Int, sizes .* 100) end -function bipartite_sc(bipartiter::KaHyParBipartite, adj::SparseMatrixCSC, vertices, log2_sizes) - n_v = length(vertices) - subgraph, remaining_edges = induced_subhypergraph(adj, vertices) - if !isdefined(@__MODULE__, :KaHyPar) - error("Module `KaHyPar` not found, please type `using KaHyPar` before using the `KaHyParBipartite` optimizer!") - end - hypergraph = KaHyPar.HyperGraph(subgraph, ones(n_v), convert2int(log2_sizes[remaining_edges])) - local parts - min_sc = 999999 - for imbalance in bipartiter.imbalances - parts = @suppress KaHyPar.partition(hypergraph, 2; imbalance=imbalance, configuration=:edge_cut) - part0 = vertices[parts .== 0] - part1 = vertices[parts .== 1] - sc0, sc1 = group_sc(adj, part0, log2_sizes), group_sc(adj, part1, log2_sizes) - sc = max(sc0, sc1) - min_sc = min(sc, min_sc) - @debug "imbalance $imbalance: sc = $sc, group = ($(length(part0)), $(length(part1)))" - if sc <= bipartiter.sc_target - return part0, part1 - end - end - error("fail to find a valid partition for `sc_target = $(bipartiter.sc_target)`, got minimum value `$min_sc` (imbalances = $(bipartiter.imbalances))") +function bipartite_sc(bipartiter, adj::SparseMatrixCSC, vertices, log2_sizes) + error("""Guess you are trying to use the `KaHyParBipartite` optimizer. +Then you need to add `using KaHyPar` first!""") end # the space complexity (external degree of freedoms) if we contract this group diff --git a/src/treesa.jl b/src/treesa.jl index fb31557..7ca226d 100644 --- a/src/treesa.jl +++ b/src/treesa.jl @@ -75,7 +75,7 @@ function Base.replace!(slicer::Slicer, pair::Pair) return slicer end -function Base.push!(slicer, best) +function Base.push!(slicer::Slicer, best) @assert length(slicer) < slicer.max_size @assert !haskey(slicer.legs, best) slicer.legs[best] = slicer.log2_sizes[best] # add best to legs diff --git a/test/interfaces.jl b/test/interfaces.jl index 6827adc..8f40e46 100644 --- a/test/interfaces.jl +++ b/test/interfaces.jl @@ -59,8 +59,8 @@ end code2 = optimize_code(code, uniformsize(code, 5), TreeSA(ntrials=1, nslices=5)) pm2 = peak_memory(code2, uniformsize(code, 5)) tc2, sc2, rw2 = timespacereadwrite_complexity(code2, uniformsize(code, 5)) - @test 5 * 2^sc1 > pm1 > 2^sc1 - @test 5 * 2^sc2 > pm2 > 2^sc2 + @test 10 * 2^sc1 > pm1 > 2^sc1 + @test 10 * 2^sc2 > pm2 > 2^sc2 end