[perf] Benchmarks and AC/AC2 contraction improvements

.gitignore

@@ -3,3 +3,4 @@ Manifest.toml
 .vscode
 .DS_Store
 **/dev/
+benchmark/results

Project.toml

@@ -53,7 +53,7 @@ RecipesBase = "1.1"
 TensorKit = "0.16.3"
 TensorKitManifolds = "0.7"
 TensorKitTensors = "0.2"
-TensorOperations = "5"
+TensorOperations = "5.5"
 Test = "1"
 TestExtras = "0.3"
 VectorInterface = "0.2, 0.3, 0.4, 0.5"

benchmark/LocalPreferences.toml

@@ -0,0 +1,2 @@
+[TensorOperations]
+precompile_workload = true

benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl

@@ -0,0 +1,59 @@
+module MPSKitBenchmarks
+
+using BenchmarkTools
+using MPSKit
+using TOML
+
+include("utils/BenchUtils.jl")
+
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 20.0
+BenchmarkTools.DEFAULT_PARAMETERS.samples = 10000
+BenchmarkTools.DEFAULT_PARAMETERS.time_tolerance = 0.15
+BenchmarkTools.DEFAULT_PARAMETERS.memory_tolerance = 0.01
+
+const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json")
+const SUITE = BenchmarkGroup()
+const MODULES = Dict{String, Symbol}(
+    "derivatives" => :DerivativesBenchmarks
+)
+
+include("derivatives/DerivativesBenchmarks.jl")
+
+load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...)
+
+function load!(group::BenchmarkGroup, id::AbstractString; tune::Bool = false)
+    modsym = MODULES[id]
+    modpath = joinpath(dirname(@__FILE__), id, "$(modsym).jl")
+    Core.eval(@__MODULE__, :(include($modpath)))
+    mod = Core.eval(@__MODULE__, modsym)
+    modsuite = @invokelatest getglobal(mod, :SUITE)
+    group[id] = modsuite
+    if tune
+        results = BenchmarkTools.load(PARAMS_PATH)[1]
+        haskey(results, id) && loadparams!(modsuite, results[id], :evals)
+    end
+    return group
+end
+
+loadall!(; kwargs...) = loadall!(SUITE; kwargs...)
+
+function loadall!(group::BenchmarkGroup; verbose::Bool = true, tune::Bool = false)
+    for id in keys(MODULES)
+        if verbose
+            print("loading group $(repr(id))... ")
+            time = @elapsed load!(group, id, tune = false)
+            println("done (took $time seconds)")
+        else
+            load!(group, id; tune = false)
+        end
+    end
+    if tune
+        results = BenchmarkTools.load(PARAMS_PATH)[1]
+        for (id, suite) in group
+            haskey(results, id) && loadparams!(suite, results[id], :evals)
+        end
+    end
+    return group
+end
+
+end

benchmark/MPSKitBenchmarks/derivatives/AC2_benchmarks.jl

@@ -0,0 +1,101 @@
+struct AC2Spec{S <: ElementarySpace} # <: BenchmarkSpec
+    physicalspaces::NTuple{2, S}
+    mps_virtualspaces::NTuple{3, S}
+    mpo_virtualspaces::NTuple{3, SumSpace{S}}
+    nonzero_keys::NTuple{2, Vector{Tuple{Int, Int}}}
+end
+
+function AC2Spec(mps, mpo; site = length(mps) ÷ 2)
+    physicalspaces = (physicalspace(mps, site), physicalspace(mps, site + 1))
+    mps_virtualspaces = (left_virtualspace(mps, site), right_virtualspace(mps, site), right_virtualspace(mps, site + 1))
+    mpo_virtualspaces = (left_virtualspace(mpo, site), right_virtualspace(mpo, site), right_virtualspace(mpo, site + 1))
+    ks = (
+        map(x -> (x.I[1], x.I[4]), nonzero_keys(mpo[site])),
+        map(x -> (x.I[1], x.I[4]), nonzero_keys(mpo[site])),
+    )
+    return AC2Spec(physicalspaces, mps_virtualspaces, mpo_virtualspaces, ks)
+end
+
+benchname(spec::AC2Spec) = dim(spec.mps_virtualspaces[1]), dim(spec.mpo_virtualspaces[1])
+
+# Benchmarks
+# ----------
+
+function prepare_state(spec::AC2Spec{S}; T::Type = Float64) where {S}
+    GL = randn(T, spec.mps_virtualspaces[1] ⊗ spec.mpo_virtualspaces[1]' ← spec.mps_virtualspaces[1])
+    GR = randn(T, spec.mps_virtualspaces[3] ⊗ spec.mpo_virtualspaces[3] ← spec.mps_virtualspaces[3])
+
+    W1 = MPSKit.JordanMPOTensor{T, S}(undef, spec.mpo_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.physicalspaces[1] ⊗ spec.mpo_virtualspaces[2])
+    for (r, c) in spec.nonzero_keys[1]
+        r == c == 1 && continue
+        r == size(W1, 1) && c == size(W1, 4) && continue
+        W1[r, 1, 1, c] = randn!(W1[r, 1, 1, c])
+    end
+    W2 = MPSKit.JordanMPOTensor{T, S}(undef, spec.mpo_virtualspaces[2] ⊗ spec.physicalspaces[2] ← spec.physicalspaces[2] ⊗ spec.mpo_virtualspaces[3])
+    for (r, c) in spec.nonzero_keys[2]
+        r == c == 1 && continue
+        r == size(W2, 1) && c == size(W2, 4) && continue
+        W2[r, 1, 1, c] = randn!(W2[r, 1, 1, c])
+    end
+    H = InfiniteMPOHamiltonian(PeriodicVector([W1, W2]))
+
+    A = PeriodicVector(
+        [
+            randn(T, spec.mps_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.mps_virtualspaces[2]),
+            randn(T, spec.mps_virtualspaces[2] ⊗ spec.physicalspaces[1] ← spec.mps_virtualspaces[3]),
+        ]
+    )
+    C = PeriodicVector(
+        [
+            randn(T, spec.mps_virtualspaces[2] ← spec.mps_virtualspaces[2]),
+            randn(T, spec.mps_virtualspaces[3] ← spec.mps_virtualspaces[3]),
+        ]
+    )
+    psi = InfiniteMPS{eltype(A), eltype(C)}(A, A, C, A)
+
+    GLs, GRs = MPSKit.initialize_environments(psi, H, psi)
+    envs = MPSKit.InfiniteEnvironments(GLs, GRs)
+
+    return psi, H, envs
+end
+
+function contraction_benchmark(spec::AC2Spec; T::Type = Float64)
+    psi, H, envs = prepare_state(spec; T)
+    H_eff = MPSKit.AC2_hamiltonian(1, psi, H, psi, envs)
+    V = spec.mps_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.mps_virtualspaces[3] ⊗ spec.physicalspaces[2]'
+    return @benchmarkable $H_eff * x setup = x = randn($T, $V)
+end
+
+function preparation_benchmark(spec::AC2Spec; T::Type = Float64)
+    psi, H, envs = prepare_state(spec; T)
+    return @benchmarkable MPSKit.AC2_hamiltonian(1, $psi, $H, $psi, $envs)
+end
+
+# Converters
+# ----------
+function tomlify(spec::AC2Spec)
+    return Dict(
+        "physicalspaces" => collect(tomlify.(spec.physicalspaces)),
+        "mps_virtualspaces" => collect(tomlify.(spec.mps_virtualspaces)),
+        "mpo_virtualspaces" => collect(tomlify.(spec.mpo_virtualspaces)),
+        "nonzero_keys" => collect(map(Base.Fix1(map, collect), spec.nonzero_keys))
+    )
+end
+
+function untomlify(::Type{AC2Spec}, x)
+    to_space = Base.Fix1(untomlify, VectorSpace)
+    physicalspaces = Tuple(map(to_space, x["physicalspaces"]))
+    mps_virtualspaces = Tuple(map(to_space, x["mps_virtualspaces"]))
+    mpo_virtualspaces = Tuple(map(to_space, x["mpo_virtualspaces"]))
+    nonzero_keys = Tuple(map(Base.Fix1(map, Tuple), x["nonzero_keys"]))
+    return AC2Spec(physicalspaces, mps_virtualspaces, mpo_virtualspaces, nonzero_keys)
+end
+
+function Base.convert(::Type{AC2Spec{S₁}}, spec::AC2Spec{S₂}) where {S₁, S₂}
+    return S₁ === S₂ ? spec : AC2Spec(
+            S₁.(spec.physicalspaces),
+            S₁.(spec.mps_virtualspaces),
+            SumSpace{S₁}.(spec.mpo_virtualspaces),
+            spec.nonzero_keys
+        )
+end

benchmark/MPSKitBenchmarks/derivatives/DerivativesBenchmarks.jl

@@ -0,0 +1,44 @@
+module DerivativesBenchmarks
+
+export AC2Spec
+
+using BenchmarkTools
+using TOML
+using TensorKit
+using BlockTensorKit
+using MPSKit
+using ..BenchUtils
+import ..BenchUtils: tomlify, untomlify
+
+const SUITE = BenchmarkGroup()
+
+const allparams = Dict(
+    "heisenberg_nn" => TOML.parsefile(joinpath(@__DIR__, "heisenberg_nn.toml")),
+    "heisenberg_nnn" => TOML.parsefile(joinpath(@__DIR__, "heisenberg_nnn.toml")),
+    "heisenberg_cylinder" => TOML.parsefile(joinpath(@__DIR__, "heisenberg_cylinder.toml")),
+    "heisenberg_coulomb" => TOML.parsefile(joinpath(@__DIR__, "heisenberg_coulomb.toml"))
+)
+
+include("AC2_benchmarks.jl")
+
+T = Float64
+
+suite_init = addgroup!(SUITE, "AC2_preparation")
+suite_apply = addgroup!(SUITE, "AC2_contraction")
+
+for (model, params) in allparams
+    g_prep = addgroup!(suite_init, model)
+    g_contract = addgroup!(suite_apply, model)
+    for (symmetry, specs) in params
+        g_prep_sym = addgroup!(g_prep, symmetry)
+        g_contract_sym = addgroup!(g_contract, symmetry)
+        for spec_dict in specs
+            spec = untomlify(AC2Spec, spec_dict)
+            name = benchname(spec)
+            g_contract_sym[name] = contraction_benchmark(spec; T)
+            g_prep_sym[name] = preparation_benchmark(spec; T)
+        end
+    end
+end
+
+end