More allocator threading

[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.

[leburgel]

This ignored the user-specified concurrency, so I replaced it with a tforeach based on the configured scheduler.

[lkdvos]

While I'm staring at this, am I seeing this wrong or does the first part of this function also cleanly separate into left- and right? In that case, it might make more sense to just have a similar code structure as in the regular infinite environments case, and split this into compute_leftenvs and compute_rightenvs functions, mirror the structure of that other part and add the timeroutput support while we are at it.

Let me know if this is straying to far from the original PR though, we can definitely also do this as a follow-up, but then we should open an issue so I don't forget

[leburgel]

I don't think it's too much work to do this here. Otherwise this will probably lose priority quite quickly anyway.

[leburgel]

I had a go at this, let me know what you think.

[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).

[lkdvos]

I think the inference issues might be because of the nested nested tuples here. However, as a different suggestion, how about something like this, and similar for the other side? This avoids the const indices, while not having to repeat the code, and still keeps everything allocated/freed within a single function so there is no leaking beyond function barriers. I think I would even manually inline the _fuse_env functions, but that is more of a personal taste thing, although if the inline annotation is still required there is not much reason not to.

function _prepare_GL_O(GL, O, backend, allocator)
    cp = allocator_checkpoint!(allocator)
    TC = TensorOperations.promote_contract(scalartype(A), scalartype(B))
    GL_O = TensorOperations.tensoralloc_contract(
        TC,
        GL, ((1, 3), (2,)), false,
        O, ((1,), (2, 3, 4)), false,
        ((1, 3, 5), (2, 4)), Val(true), allocator
    )
    @plansor backend = backend allocator = allocator begin
        GL_O[-1 -2 -3; -4 -5] = GL[-1 1; -4] * O[1 -2; -5 -3]
    end
    left_env = _fuse_env(GL_O, 1, 2, backend, allocator)
    
    TensorOperations.tensor_free!(GL_O, allocator)
    allocator_reset!(allocator, cp)
    
    return left_env
end

[leburgel]

I liked the suggestion so I rewrote it like this. _fuse_env still needs an @inline to guarantee the inference test, but since it's used in both preparations it seemed best to keep them separate?

[lkdvos]

Yeah, the inline is because it requires constprop on the N1, N2, so it has to either be Val-based or just inlined. I don't think it matters that much, so just accepted and merged this :)

[lkdvos]

While I'm staring at this, am I seeing this wrong or does the first part of this function also cleanly separate into left- and right? In that case, it might make more sense to just have a similar code structure as in the regular infinite environments case, and split this into compute_leftenvs and compute_rightenvs functions, mirror the structure of that other part and add the timeroutput support while we are at it.

Let me know if this is straying to far from the original PR though, we can definitely also do this as a follow-up, but then we should open an issue so I don't forget