TensorBFS / OMEinsumContractionOrders.jl

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

dump load contraction orders #24

Closed GiggleLiu closed 2 years ago

GiggleLiu commented 2 years ago
julia> using OMEinsum, OMEinsumContractionOrders, Graphs
[ Info: Precompiling OMEinsumContractionOrders [6f22d1fd-8eed-4bb7-9776-e7d684900715]

julia> function random_regular_eincode(n, k; optimize=nothing)
            g = Graphs.random_regular_graph(n, k)
            ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
            return EinCode((ixs..., [(i,) for i in     Graphs.vertices(g)]...), ())
           end
random_regular_eincode (generic function with 1 method)

julia> code = random_regular_eincode(200, 3);

julia> optcode_tree = optimize_code(code, uniformsize(code, 2), TreeSA(sc_target=28, βs=0.1:0.1:20,
                                                                ntrials=5, niters=30, sc_weight=2.0));

julia> timespace_complexity(optcode_tree, uniformsize(code, 2))
(30.536222605629487, 26.0)

julia> writejson("test.json", optcode_tree)

julia> optcode_loaded = readjson("test.json")
SlicedEinsum{Any, NestedEinsum{DynamicEinCode{Int64}}}(Slicing{Any}(Any[]), 181, 181 -> 
├─ 145∘181, 181∘145 -> 181
│  ├─ 75∘181∘26, 145∘75∘26 -> 181∘145
│  │  ├─ 145∘191∘90, 145∘75∘26∘191∘90 -> 145∘75∘26
│  │  │  ├─ 82∘96∘176∘47∘37∘145, 47∘176∘75∘82∘26∘37∘191∘90∘96 -> 145∘75∘26∘191∘90
│  │  │  │  ├─ 90∘70∘75∘47∘105∘96∘72∘163∘200∘150∘183∘66∘196∘85∘50∘191, 196∘105∘150∘50∘47∘183∘163∘176∘70∘75∘85∘82∘72∘66∘26∘200∘37∘191 -> 47∘176∘75∘82∘26∘37∘191∘90∘96
│  │  │  │  │  ⋮
│  │  │  │  │  
│  │  │  │  └─ 82∘168∘96∘176, 168∘47∘37∘145 -> 82∘96∘176∘47∘37∘145
│  │  │  │     ⋮
│  │  │  │     
│  │  │  └─ 90∘107, 107∘145∘191 -> 145∘191∘90
│  │  │     ├─ 107∘191, 107∘145 -> 107∘145∘191
│  │  │     │  ⋮
│  │  │     │  
│  │  │     └─ 90∘107
│  │  └─ 26∘181, 75∘181 -> 75∘181∘26
│  │     ├─ 75∘181
│  │     └─ 26∘181
│  └─ 145, 145∘181 -> 145∘181
│     ├─ 145∘181
│     └─ 145
└─ 181
)

julia> timespace_complexity(optcode_loaded, uniformsize(code, 2))
(30.536222605629487, 26.0)

Data format

julia> code = EinCode([[1,2], [2,3], [3,4]], [1,4])
1∘2, 2∘3, 3∘4 -> 1∘4

julia> optcode = optimize_code(code, uniformsize(code, 2), TreeSA(; nslices=1))
SlicedEinsum{Int64, NestedEinsum{DynamicEinCode{Int64}}}(Slicing{Int64}([2]), 2∘4, 1∘2 -> 1∘4
├─ 1∘2
└─ 3∘4, 2∘3 -> 2∘4
   ├─ 2∘3
   └─ 3∘4
)

Will generate file

{
    "slices": [
        2
    ],
    "output": [
        1,
        4
    ],
    "tree": {
        "isleaf": false,
        "eins": {
            "ixs": [
                [
                    2,
                    4
                ],
                [
                    1,
                    2
                ]
            ],
            "iy": [
                1,
                4
            ]
        },
        "args": [
            {
                "isleaf": false,
                "eins": {
                    "ixs": [
                        [
                            3,
                            4
                        ],
                        [
                            2,
                            3
                        ]
                    ],
                    "iy": [
                        2,
                        4
                    ]
                },
                "args": [
                    {
                        "isleaf": true,
                        "tensorindex": 3
                    },
                    {
                        "isleaf": true,
                        "tensorindex": 2
                    }
                ]
            },
            {
                "isleaf": true,
                "tensorindex": 1
            }
        ]
    },
    "label-type": "Int64",
    "inputs": [
        [
            1,
            2
        ],
        [
            2,
            3
        ],
        [
            3,
            4
        ]
    ]
}

julia> @benchmark CUDA.@sync optcode_loaded([CUDA.rand(fill(2, length(ix))...) for ix in getixsv(optcode_loaded)]...) BenchmarkTools.Trial: 12 samples with 1 evaluation. Range (min … max): 417.572 ms … 485.062 ms ┊ GC (min … max): 2.28% … 4.11% Time (median): 441.641 ms ┊ GC (median): 4.45% Time (mean ± σ): 444.844 ms ± 24.147 ms ┊ GC (mean ± σ): 3.69% ± 1.22%

█ █ ██ █ █ █ █ █ █ █ █
█▁█▁██▁▁█▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁█▁█▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁█ ▁ 418 ms Histogram: frequency by time 485 ms <

Memory estimate: 42.38 MiB, allocs estimate: 739486.

codecov[bot] commented 2 years ago

Codecov Report

Merging #24 (35e8eab) into master (8903b36) will increase coverage by 0.29%. The diff coverage is 100.00%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master      #24      +/-   ##
==========================================
+ Coverage   95.46%   95.75%   +0.29%     
==========================================
  Files           7        8       +1     
  Lines         684      731      +47     
==========================================
+ Hits          653      700      +47     
  Misses         31       31              
Impacted Files Coverage Δ
src/OMEinsumContractionOrders.jl 100.00% <ø> (ø)
src/json.jl 100.00% <100.00%> (ø)

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 8903b36...35e8eab. Read the comment docs.