use backends and allocators in DMRG

Project.toml

@@ -18,6 +18,7 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
@@ -46,12 +47,13 @@ LoggingExtras = "~1.0"
 MatrixAlgebraKit = "0.6.5"
 OhMyThreads = "0.7, 0.8"
 OptimKit = "0.3.1, 0.4, 0.5"
+Preferences = "1"
 Printf = "1"
 Random = "1"
 RecipesBase = "1.1"
 TensorKit = "0.17.1"
 TensorKitManifolds = "0.7, 0.8"
-TensorOperations = "5.5.1"
+TensorOperations = "5.6.2"
 TimerOutputs = "0.5.29"
 VectorInterface = "0.2, 0.3, 0.4, 0.5, 0.6"
 julia = "1.10"

benchmark/gpu-allocator/run_gpu_tests.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+# Run the MPSKit CUDA test suite, including the new algorithm-level allocator tests.
+#
+# There is no GPU on the workstation, so this is the only way to exercise
+# `MPSKit.allocatortype(::Type{<:CuArray})` and confirm that a device-backed DMRG/TDVP draws its
+# scratch space from `CUDAAllocator` rather than a host buffer.
+#
+# Submit from the repository root:
+#     sbatch benchmark/gpu-allocator/run_gpu_tests.sbatch
+#
+# Adjust --partition/--gres to whatever GPU partition you have access to; the values below are
+# placeholders. Check with `sinfo -o "%P %G %N"`.
+
+#SBATCH --job-name=mpskit-gpu-alloc
+#SBATCH --output=benchmark/gpu-allocator/logs/%x-%j.out
+#SBATCH --error=benchmark/gpu-allocator/logs/%x-%j.err
+#SBATCH --time=01:00:00
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=64G
+#SBATCH --partition=gpu          # placeholder
+#SBATCH --gres=gpu:1             # placeholder
+
+set -euo pipefail
+
+cd "${SLURM_SUBMIT_DIR:?must be submitted with sbatch}"
+mkdir -p benchmark/gpu-allocator/logs
+
+# The depot is precompiled under 1.12.6; `module load julia` gives 1.11.2, which fails with a
+# StaticData error. Use the juliaup install on GPFS directly.
+JULIA="${HOME}/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/bin/julia"
+test -x "$JULIA" || { echo "julia not found at $JULIA" >&2; exit 1; }
+
+echo "host:     $(hostname)"
+echo "julia:    $("$JULIA" --version)"
+nvidia-smi || echo "WARNING: nvidia-smi unavailable"
+
+export JULIA_NUM_THREADS=4
+# one BLAS thread per Julia thread would oversubscribe; the GPU does the heavy lifting anyway
+export OPENBLAS_NUM_THREADS=1
+
+"$JULIA" --project=test -e '
+    using Pkg
+    Pkg.instantiate()
+
+    using CUDA, cuTENSOR
+    @info "CUDA status" functional = CUDA.functional() cutensor = cuTENSOR.functional()
+    CUDA.functional() || error("CUDA is not functional on this node")
+    CUDA.versioninfo()
+
+    include(joinpath(@__DIR__, "test", "setup", "testsetup.jl"))
+    using .TestSetup
+    using Test
+
+    @testset "MPSKit CUDA" verbose = true begin
+        for f in ("cu_adapt.jl", "states.jl", "operators.jl", "algorithms.jl")
+            @testset "$f" begin
+                include(joinpath(@__DIR__, "test", "gpu", "cuda", f))
+            end
+        end
+    end
+'

docs/src/man/parallelism.md

@@ -98,6 +98,29 @@ MPSKit.Defaults.set_scheduler!(:dynamic) # default: multithreading with some loa
 For further reference on the available schedulers and finer control, please refer to the
 [`OhMyThreads.jl` documentation](https://juliafolds2.github.io/OhMyThreads.jl/stable/)
 
+## Allocators
+
+Every local update in an MPS algorithm contracts a handful of intermediate tensors that are thrown away immediately afterwards.
+Rather than asking Julia's garbage collector to clean up after each one, MPSKit hands the contractions an allocator that manages those intermediates itself.
+On a typical ground-state search this cuts both the allocation count and the garbage-collection time substantially.
+
+Which allocator is appropriate depends on two things, so [`MPSKit.default_allocator`](@ref) is asked for one at the start of every solve, and the answer is then used for all of its local updates:
+
+- *Where the tensors live.* Host memory can be served by an allocator that bypasses Julia's memory manager, whereas a device-backed state needs its scratch space on the device. This is only known once there is a state in hand, which is why it is not a setting on the algorithm.
+- *Whether the allocator will be shared between tasks.* `TensorOperations.BufferAllocator` is a bump buffer with a mutable offset, so it is only safe when a single task owns it - a serial sweep. Concurrent work instead gets a `TensorOperations.ManualAllocator`, which `malloc`s and `free`s each intermediate and holds no state at all.
+
+Anything MPSKit does not recognise falls back on `TensorOperations.DefaultAllocator`, which allocates through the storage type itself: correct on any device, at the cost of leaving the intermediates to the garbage collector.
+That is also what device-backed states get.
+Note that the `backend` field needs no such treatment: `TensorOperations.DefaultBackend` is a placeholder that resolves to an implementation based on the types of the tensors involved, so it already selects cuTENSOR for `CuArray`-backed states.
+
+Dedicated scratch space can be turned off when memory rather than time is the binding constraint, which trades it back for garbage-collector pressure:
+
+```julia
+MPSKit.Defaults.set_buffering!(false)
+```
+
+This is a compile-time preference: it is written to a `LocalPreferences.toml` file next to the active `Project.toml`, and Julia has to be restarted for the change to take effect.
+
 ## TensorKit multithreading
 
 Finally, when dealing with tensors that have some internal symmetry, it is also possible to

src/MPSKit.jl

@@ -56,6 +56,7 @@ using Compat: @compat
 @compat public DynamicTols, DynamicTol, AdaptiveKrylov
 @compat public VERBOSE_NONE, VERBOSE_WARN, VERBOSE_CONV, VERBOSE_ITER, VERBOSE_ALL
 @compat public IterLog, loginit!, logiter!, logfinish!, logcancel!
+@compat public default_allocator
 
 # Imports
 # -------
@@ -68,7 +69,7 @@ using BlockTensorKit
 using BlockTensorKit: TensorMapSumSpace
 using TensorOperations
 using TensorOperations: AbstractBackend, DefaultBackend, DefaultAllocator, BufferAllocator,
-    allocator_checkpoint!, allocator_reset!
+    ManualAllocator, allocator_checkpoint!, allocator_reset!
 using KrylovKit
 using KrylovKit: KrylovAlgorithm
 using OptimKit
@@ -98,6 +99,7 @@ using .DynamicTols
 
 include("utility/defaults.jl")
 using .Defaults: VERBOSE_NONE, VERBOSE_WARN, VERBOSE_CONV, VERBOSE_ITER, VERBOSE_ALL
+include("utility/allocator.jl")
 include("utility/logging.jl")
 using .IterativeLoggers
 include("utility/iterativesolvers.jl")

src/algorithms/ED.jl

@@ -2,7 +2,8 @@
     exact_diagonalization(H::FiniteMPOHamiltonian;
                           sector=rightunit(H),
                           len::Int=length(H), num::Int=1, which::Symbol=:SR,
-                          alg=Defaults.alg_eigsolve(; dynamic_tols=false))
+                          alg=Defaults.alg_eigsolve(; dynamic_tols=false),
+                          backend=Defaults.backend())
                             -> vals, state_vecs, convhist
 
 Use [`KrylovKit.eigsolve`](@extref) to perform exact diagonalization on a
@@ -20,6 +21,7 @@ equivalent to dense eigenvectors.
 - `which::Symbol=:SR`: the kind eigenvalues to find, see [`KrylovKit.eigsolve`](@extref). 
 - `alg=Defaults.alg_eigsolve(; dynamic_tols=false)`: the diagonalization algorithm to use,
   see [`KrylovKit.eigsolve`](@extref).
+- `backend=DefaultBackend()`: backend for tensor contractions and index manipulations.
 
 !!! note "Valid `sector` values"
     The total charge of the eigenvectors is imposed by adding a charged auxiliary space as
@@ -33,7 +35,8 @@ equivalent to dense eigenvectors.
 function exact_diagonalization(
         H::FiniteMPOHamiltonian;
         sector = rightunit(H), num::Int = 1, which::Symbol = :SR,
-        alg = Defaults.alg_eigsolve(; dynamic_tols = false)
+        alg = Defaults.alg_eigsolve(; dynamic_tols = false),
+        backend::AbstractBackend = Defaults.backend()
     )
     L = length(H)
     @assert L > 1 "FiniteMPOHamiltonian must have length > 1"
@@ -73,7 +76,8 @@ function exact_diagonalization(
     # optimize the middle site
     # Because the MPS is full rank - this is equivalent to the full Hamiltonian
     AC₀ = state.AC[middle_site]
-    H_ac = AC_hamiltonian(middle_site, state, H, state, envs)
+    allocator = default_allocator(state, SerialScheduler())
+    H_ac = AC_hamiltonian(middle_site, state, H, state, envs; backend, allocator)
     vals, vecs, convhist = eigsolve(H_ac, AC₀, num, which, alg)
 
     # repack eigenstates

src/algorithms/approximate/fvomps.jl

@@ -1,4 +1,5 @@
 function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environments(ψ, _environment_args(Oϕ)...))
+    allocator = default_allocator(ψ, SerialScheduler())
     ϵ::Float64 = 2 * alg.tol
     log = IterLog("DMRG2")
 
@@ -7,7 +8,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environment
         for iter in 1:(alg.maxiter)
             ϵ = 0.0
             for pos in [1:(length(ψ) - 1); (length(ψ) - 2):-1:1]
-                AC2′ = AC2_projection(pos, ψ, Oϕ, envs)
+                AC2′ = AC2_projection(pos, ψ, Oϕ, envs; alg.backend, allocator)
                 al, c, ar, = svd_trunc!(AC2′, inner_alg_gauge(alg))
 
                 AC2 = ψ.AC[pos] * _transpose_tail(ψ.AR[pos + 1])
@@ -36,6 +37,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environment
 end
 
 function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG, envs = environments(ψ, _environment_args(Oϕ)...))
+    allocator = default_allocator(ψ, SerialScheduler())
     ϵ::Float64 = 2 * alg.tol
     log = IterLog("DMRG")
 
@@ -44,7 +46,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG, envs = environments
         for iter in 1:(alg.maxiter)
             ϵ = 0.0
             for pos in [1:(length(ψ) - 1); length(ψ):-1:2]
-                AC′ = AC_projection(pos, ψ, Oϕ, envs)
+                AC′ = AC_projection(pos, ψ, Oϕ, envs; alg.backend, allocator)
                 AC = ψ.AC[pos]
                 ϵ = max(ϵ, norm(AC′ - AC) / norm(AC′))
 

src/algorithms/approximate/idmrg.jl

@@ -2,6 +2,7 @@ function approximate!(
         ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS}, alg::IDMRG,
         envs = environments(ψ, toapprox...)
     )
+    allocator = default_allocator(ψ, SerialScheduler())
     log = IterLog("IDMRG")
     ϵ::Float64 = 2 * alg.tol
     local iter
@@ -15,7 +16,8 @@ function approximate!(
             for col in 1:size(ψ, 2)
                 for row in 1:size(ψ, 1)
                     ψ.AC[row + 1, col] = AC_projection(
-                        CartesianIndex(row, col), ψ, toapprox, envs
+                        CartesianIndex(row, col), ψ, toapprox, envs;
+                        alg.backend, allocator
                     )
                     normalize!(ψ.AC[row + 1, col])
                     ψ.AL[row + 1, col], ψ.C[row + 1, col] = left_orth!(ψ.AC[row + 1, col])
@@ -27,7 +29,8 @@ function approximate!(
             for col in reverse(1:size(ψ, 2))
                 for row in 1:size(ψ, 1)
                     ψ.AC[row + 1, col] = AC_projection(
-                        CartesianIndex(row, col), ψ, toapprox, envs
+                        CartesianIndex(row, col), ψ, toapprox, envs;
+                        alg.backend, allocator
                     )
                     normalize!(ψ.AC[row + 1, col])
                     ψ.C[row + 1, col - 1], temp = right_orth!(_transpose_tail(ψ.AC[row + 1, col]))
@@ -64,6 +67,7 @@ function approximate!(
         ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS},
         alg::IDMRG2, envs = environments(ψ, toapprox...)
     )
+    allocator = default_allocator(ψ, SerialScheduler())
     size(ψ, 2) < 2 && throw(ArgumentError("unit cell should be >= 2"))
     ϵ::Float64 = 2 * alg.tol
     log = IterLog("IDMRG2")
@@ -80,7 +84,7 @@ function approximate!(
                 for row in 1:size(ψ, 1)
                     AC2′ = AC2_projection(
                         CartesianIndex(row, site), ψ, toapprox, envs;
-                        kind = :ACAR
+                        kind = :ACAR, alg.backend, allocator
                     )
                     al, c, ar = svd_trunc!(AC2′; trunc = alg.trscheme, alg = alg.alg_svd)
                     normalize!(c)
@@ -99,7 +103,10 @@ function approximate!(
             ψ.AL[1, end] = ψ.AC[1, end] / ψ.C[1, end]
             ψ.AC[1, 1] = _mul_tail(ψ.AL[1, 1], ψ.C[1, 1])
             for row in 1:size(ψ, 1)
-                AC2′ = AC2_projection(CartesianIndex(row, size(ψ, 2)), ψ, toapprox, envs; kind = :ALAC)
+                AC2′ = AC2_projection(
+                    CartesianIndex(row, size(ψ, 2)), ψ, toapprox, envs;
+                    kind = :ALAC, alg.backend, allocator
+                )
                 al, c, ar = svd_trunc!(AC2′; trunc = alg.trscheme, alg = alg.alg_svd)
                 normalize!(c)
 
@@ -123,7 +130,7 @@ function approximate!(
                 for row in 1:size(ψ, 1)
                     AC2′ = AC2_projection(
                         CartesianIndex(row, site), ψ, toapprox, envs;
-                        kind = :ALAC
+                        kind = :ALAC, alg.backend, allocator
                     )
                     al, c, ar = svd_trunc!(AC2′; trunc = alg.trscheme, alg = alg.alg_svd)
                     normalize!(c)
@@ -141,7 +148,10 @@ function approximate!(
             ψ.AC[1, end] = _mul_front(ψ.C[1, end - 1], ψ.AR[1, end])
             ψ.AR[1, 1] = _transpose_front(ψ.C[1, end] \ _transpose_tail(ψ.AC[1, 1]))
             for row in 1:size(ψ, 1)
-                AC2′ = AC2_projection(CartesianIndex(row, 0), ψ, toapprox, envs; kind = :ACAR)
+                AC2′ = AC2_projection(
+                    CartesianIndex(row, 0), ψ, toapprox, envs;
+                    kind = :ACAR, alg.backend, allocator
+                )
                 al, c, ar = svd_trunc!(AC2′; trunc = alg.trscheme, alg = alg.alg_svd)
                 normalize!(c)
 
@@ -155,7 +165,6 @@ function approximate!(
             transfer_leftenv!(envs, ψ, toapprox, 1)
             transfer_rightenv!(envs, ψ, toapprox, 0)
 
-
             normalize!(envs, ψ, toapprox)
 
             # update error

src/algorithms/approximate/vomps.jl

@@ -5,7 +5,10 @@ Base.@deprecate(
     ),
     approximate(
         ψ, toapprox,
-        VOMPS(; alg.tol, alg.maxiter, alg.finalize, alg.verbosity, alg.alg_gauge, alg.alg_environments),
+        VOMPS(;
+            alg.tol, alg.maxiter, alg.finalize, alg.verbosity, alg.alg_gauge,
+            alg.alg_environments, alg.backend,
+        ),
         envs...; kwargs...
     )
 )
@@ -14,6 +17,10 @@ function approximate(
         mps::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS}, alg::VOMPS,
         envs = environments(mps, toapprox...)
     )
+    return _approximate_vomps(mps, toapprox, alg, envs)
+end
+
+function _approximate_vomps(mps, toapprox, alg::VOMPS, envs)
     log = IterLog("VOMPS")
     iter = 0
     ϵ = calc_galerkin(mps, toapprox..., envs)
@@ -71,17 +78,24 @@ function localupdate_step!(
     ACs = similar(state.mps.AC)
     dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
 
-    foreach(eachsite(state.mps)) do site
+    # the sweep is serial, so a single allocator serves all sites
+    allocator = default_allocator(state.mps, SerialScheduler())
+    for site in eachsite(state.mps)
         AC = map(1:size(state.mps, 1)) do row
-            AC_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs)
+            AC_projection(
+                CartesianIndex(row, site), state.mps, state.operator, state.envs;
+                it.backend, allocator
+            )
         end
         circshift!(AC, 1)
         C = map(1:size(state.mps, 1)) do row
-            C_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs)
+            C_projection(
+                CartesianIndex(row, site), state.mps, state.operator, state.envs;
+                it.backend, allocator
+            )
         end
         circshift!(C, 1)
         dst_ACs[site] = regauge!(AC, C; alg = alg_orth)
-        return nothing
     end
 
     return ACs
@@ -95,18 +109,27 @@ function localupdate_step!(
     ACs = similar(state.mps.AC)
     dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
 
+    # every site - and the AC and C projections within a site - runs concurrently, so the allocator
+    # is shared and has to be one that tolerates that
+    allocator = default_allocator(state.mps, scheduler)
     tforeach(eachsite(state.mps); scheduler) do site
         local AC, C
         @sync begin
             Threads.@spawn begin
                 AC = map(1:size(state.mps, 1)) do row
-                    AC_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs)
+                    AC_projection(
+                        CartesianIndex(row, site), state.mps, state.operator, state.envs;
+                        it.backend, allocator
+                    )
                 end
                 circshift!(AC, 1)
             end
             Threads.@spawn begin
                 C = map(1:size(state.mps, 1)) do row
-                    C_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs)
+                    C_projection(
+                        CartesianIndex(row, site), state.mps, state.operator, state.envs;
+                        it.backend, allocator
+                    )
                 end
                 circshift!(C, 1)
             end