[leburgel]

Extends the existing backend and allocator handling to thread these through to more subroutines:

• transfer matrices and infinite environments, which makes the biggest impact since this is still about half the cost in VUMPS
• gauging sweeps
• finite environments
• quasiparticle excitations

I have some bencharks on the way, but from the partial results I've seen this gives quite a big additional reduction in allocations on top of the initial allocator effort.

[codecov]

Codecov Report

❌ Patch coverage is 88.55219% with 34 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/environments/qp_envs.jl	88.23%	12 Missing ⚠️
src/transfermatrix/transfer.jl	68.42%	12 Missing ⚠️
src/environments/multiline_envs.jl	63.63%	4 Missing ⚠️
src/algorithms/derivatives/mpo_derivatives.jl	92.30%	3 Missing ⚠️
src/states/ortho.jl	89.47%	2 Missing ⚠️
...c/algorithms/excitation/quasiparticleexcitation.jl	94.44%	1 Missing ⚠️

Files with missing lines	Coverage Δ
.../algorithms/derivatives/hamiltonian_derivatives.jl	89.08% <100.00%> (ø)
src/algorithms/excitation/exci_transfer_system.jl	96.15% <100.00%> (ø)
src/algorithms/groundstate/idmrg.jl	99.36% <100.00%> (ø)
src/algorithms/groundstate/vumps.jl	100.00% <100.00%> (+1.33%)	⬆️
src/environments/finite_envs.jl	94.73% <100.00%> (ø)
src/environments/infinite_envs.jl	95.18% <100.00%> (-0.06%)	⬇️
src/environments/multiple_envs.jl	78.37% <100.00%> (+5.40%)	⬆️
src/transfermatrix/transfermatrix.jl	80.00% <100.00%> (ø)
...c/algorithms/excitation/quasiparticleexcitation.jl	80.28% <94.44%> (+0.28%)	⬆️
src/states/ortho.jl	86.61% <89.47%> (ø)
... and 4 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[leburgel]

Some benchmark results:

Benchmark script

benchmark/allocator_benchmark.jl

# Allocation benchmark for the backend/allocator threading.
#
#     julia --project=benchmark benchmark/allocator_benchmark.jl default
#     julia --project=benchmark benchmark/allocator_benchmark.jl buffer
#
# One allocator per process: selecting it is a method redefinition of `default_allocator`, so
# it cannot be varied within a run. Invoke once per allocator and compare the two outputs;
# to measure a change, do the same on both revisions.
#
# Every case is converged first and measured separately, so compilation is not charged to the
# measurement, and the solver tolerance is pinned to zero so each case runs exactly `ITERS`
# iterations -- with a finite tolerance a nearly-converged case exits early while the reported
# figure still divides by `ITERS`, which makes it meaningless.
#
# BLAS is held to one thread so the figure is about allocation rather than threading.
#
# Environment variables: BENCH_ITERS (default 10), BENCH_CHI (default 256),
# BENCH_CYLINDER_CHI (default 128), BENCH_CASES (comma-separated subset of the case names).

using LinearAlgebra, Printf, Random
using TensorOperations: BufferAllocator, DefaultAllocator
using MPSKit, MPSKitModels, TensorKit

BLAS.set_num_threads(1)

const KIND = length(ARGS) >= 1 ? lowercase(ARGS[1]) : "buffer"
KIND in ("default", "buffer") || error("argument must be `default` or `buffer`")
const ITERS = parse(Int, get(ENV, "BENCH_ITERS", "10"))
const CHI = parse(Int, get(ENV, "BENCH_CHI", "256"))
const CYL_CHI = parse(Int, get(ENV, "BENCH_CYLINDER_CHI", "128"))

let alloc = KIND == "default" ? DefaultAllocator() : BufferAllocator()
    @eval MPSKit default_allocator(::Type{<:MPSKit.HostStorage}, ::MPSKit.SerialScheduler) = $alloc
end

# Models
# ------
# `tol = 0` everywhere: see the note above.
gs_alg(which, chi, n) = which === :idmrg2 ?
    IDMRG2(; trunc = truncrank(chi), maxiter = n, tol = 0.0, verbosity = 0) :
    VUMPS(; maxiter = n, tol = 0.0, verbosity = 0)

"Quantum nearest-neighbour Ising, 2-site cell. Dense, no symmetries: the cheap case."
function case_ising_quantum(which)
    H = repeat(transverse_field_ising(; g = 1.5), 2)
    P = physicalspace(H, 1)
    Random.seed!(1234)
    psi0 = InfiniteMPS([P, P], [typeof(P)(CHI), typeof(P)(CHI)])
    psi, = find_groundstate(psi0, H, gs_alg(which, CHI, 5))
    return () -> find_groundstate(psi, H, gs_alg(which, CHI, ITERS), environments(psi, H, psi))
end

"""
Heisenberg S=1 on a width-6 infinite cylinder — the same case MPSKit's own
`benchmark/MPSKitBenchmarks/scripts/setup_heisenberg.jl` uses as its heavy example. Wide MPO,
`BlockTensorMap` environments: this is where the allocator matters.

Run for both symmetries, because they stress different things: `SU2Irrep` gives many small
symmetry blocks, `Trivial` gives one large dense block at the same physical bond dimension.
"""
function case_heisenberg_cylinder(which, symmetry = SU2Irrep)
    symmetry === Trivial && return case_heisenberg_cylinder_trivial(which)
    H = heisenberg_XXX(Float64, SU2Irrep, InfiniteCylinder(6); spin = 1)
    Random.seed!(1234)
    # Virtual spaces taken from MPSKit's own cylinder benchmark
    # (`benchmark/MPSKitBenchmarks/derivatives/heisenberg_cylinder.toml`). Integer SU(2)
    # charges, not half-integer: a six-site spin-1 unit cell fuses to integer total spin, and a
    # trivial or half-integer virtual space admits no fusion channels at all -- the initial
    # tensors come out rank-zero and the first eigensolve throws `initial vector should not
    # have norm zero`.
    V = Rep[SU₂](0 => 1, 1 => 1, 2 => 1, 3 => 1)
    psi0 = InfiniteMPS(collect(physicalspace(H, i) for i in 1:length(H)), fill(V, length(H)))
    # Grow into a nontrivial bond dimension before measuring: a product state measures nothing.
    psi, = find_groundstate(psi0, H, gs_alg(:idmrg2, CYL_CHI, 8))
    return () -> find_groundstate(psi, H, gs_alg(which, CYL_CHI, ITERS), environments(psi, H, psi))
end

function case_heisenberg_cylinder_trivial(which)
    H = heisenberg_XXX(Float64, Trivial, InfiniteCylinder(6); spin = 1)
    Random.seed!(1234)
    # Same benchmark file, `Trivial` entries: physical ℂ^3, virtual ℂ^3 at the small end.
    V = ComplexSpace(8)
    psi0 = InfiniteMPS(collect(physicalspace(H, i) for i in 1:length(H)), fill(V, length(H)))
    psi, = find_groundstate(psi0, H, gs_alg(:idmrg2, CYL_CHI, 8))
    return () -> find_groundstate(psi, H, gs_alg(which, CYL_CHI, ITERS), environments(psi, H, psi))
end

"Classical 2D Ising transfer matrix: a plain `InfiniteMPO` through `leading_boundary`."
function case_ising_classical(which)
    O = repeat(classical_ising(; beta = log(1 + sqrt(2)) / 2), 2)   # 2-site: IDMRG2 needs >= 2
    P = physicalspace(O, 1)
    Random.seed!(1234)
    psi0 = InfiniteMPS([P, P], [typeof(P)(CHI), typeof(P)(CHI)])
    psi, = leading_boundary(psi0, O, gs_alg(which, CHI, 5))
    return () -> leading_boundary(psi, O, gs_alg(which, CHI, ITERS), environments(psi, O, psi))
end

"""
Quasiparticle excitations. `excitations` re-applies the effective excitation Hamiltonian on
every Krylov iteration and each application rebuilds the quasiparticle environments, so this
is a per-iteration allocation path. Reported per solve rather than per iteration.
"""
function case_excitations(_)
    H = repeat(transverse_field_ising(; g = 1.5), 2)
    P = physicalspace(H, 1)
    Random.seed!(1234)
    psi0 = InfiniteMPS([P, P], [typeof(P)(CHI), typeof(P)(CHI)])
    psi, envs = find_groundstate(psi0, H, VUMPS(; maxiter = 25, tol = 1.0e-10, verbosity = 0))
    excitations(H, QuasiparticleAnsatz(), 0.0, psi, envs; num = 1)   # warm up
    return () -> excitations(H, QuasiparticleAnsatz(), 0.0, psi, envs; num = 1)
end

const CASES = [
    ("ising-quantum  idmrg2", () -> case_ising_quantum(:idmrg2), ITERS),
    ("ising-quantum  vumps ", () -> case_ising_quantum(:vumps), ITERS),
    ("heisenberg-cyl-su2 idmrg2", () -> case_heisenberg_cylinder(:idmrg2, SU2Irrep), ITERS),
    ("heisenberg-cyl-su2 vumps ", () -> case_heisenberg_cylinder(:vumps, SU2Irrep), ITERS),
    ("heisenberg-cyl-triv idmrg2", () -> case_heisenberg_cylinder(:idmrg2, Trivial), ITERS),
    ("heisenberg-cyl-triv vumps ", () -> case_heisenberg_cylinder(:vumps, Trivial), ITERS),
    ("ising-classical idmrg2", () -> case_ising_classical(:idmrg2), ITERS),
    ("ising-classical vumps ", () -> case_ising_classical(:vumps), ITERS),
    ("excitations    qp    ", () -> case_excitations(nothing), 1),
]

selected = get(ENV, "BENCH_CASES", "")
cases = isempty(selected) ? CASES :
    filter(c -> any(occursin(strip(s), first(c)) for s in split(selected, ",")), CASES)

@printf("allocator = %s,  iters = %d,  chi = %d (cylinder %d),  BLAS threads = 1\n\n",
        KIND, ITERS, CHI, CYL_CHI)
@printf("%-24s %10s %14s %14s %8s\n", "case", "time [s]", "total [GiB]", "MiB/unit", "GC")
println("-"^76)
for (name, setup, unit) in cases
    run = try
        setup()
    catch e
        @printf("%-24s  SETUP FAILED: %s\n", name, first(sprint(showerror, e), 80))
        continue
    end
    GC.gc(true)
    stats = try
        @timed run()
    catch e
        @printf("%-24s  RUN FAILED: %s\n", name, first(sprint(showerror, e), 80))
        continue
    end
    @printf("%-24s %10.2f %14.3f %14.2f %7.1f%%\n", name, stats.time,
            stats.bytes / 2^30, stats.bytes / 2^20 / unit, 100 * stats.gctime / stats.time)
end

Results (MiB allocated per iteration):

case	before / default	before / buffer	after / default	after / buffer	this PR
quantum Ising, IDMRG2	2215.2	1229.8	2253.1	895.1	−27 %
quantum Ising, VUMPS	31924.2	31738.8	31924.1	12715.2	−60 %
Heisenberg cylinder SU(2), IDMRG2	676.1	405.8	705.7	385.4	−5 %
Heisenberg cylinder SU(2), VUMPS	361.9	320.8	362.6	231.4	−28 %
Heisenberg cylinder Trivial, IDMRG2	29943.6	7090.6	29967.0	3912.4	−45 %
Heisenberg cylinder Trivial, VUMPS	11822.1	8360.0	11822.2	2489.4	−70 %
classical Ising, IDMRG2	2916.2	2518.2	2916.6	1283.4	−49 %
classical Ising, VUMPS	14916.2	13912.2	14916.2	6301.2	−55 %
quasiparticle excitations	—	212984.5	212984.4	98650.6	−54 %

GC time over the same comparison: 11.3 % → 6.7 % (quantum Ising VUMPS), 15.4 % → 8.3 %
(Trivial cylinder VUMPS), 15.3 % → 11.7 % (classical Ising VUMPS), 19.2 % → 15.0 % (excitations).