[WIP] GPU support in PEPSKit

.buildkite/pipeline.yml

@@ -23,7 +23,7 @@ steps:
           - "1.13"
         queue:
           - "cuda"
-          #- "rocm" # disable for now as these are timing out 
+          - "rocm"
         group:
           - "boundarymps"
           - "bondenv"

Project.toml

@@ -33,18 +33,21 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
+[sources]
+TensorKit = {rev = "main", url = "https://github.com/QuantumKitHub/TensorKit.jl"}
+
 [extensions]
 PEPSKitAdaptExt = "Adapt"
 PEPSKitGPUArraysExt = "GPUArrays"
 
 [compat]
 Accessors = "0.1"
 Adapt = "4"
+GPUArrays = "11, 12"
 ChainRulesCore = "1.0"
 Compat = "3.46, 4.2"
 DocStringExtensions = "0.9.3"
 FiniteDifferences = "0.12"
-GPUArrays = "11, 12"
 KrylovKit = "0.9.5, 0.10"
 LinearAlgebra = "1"
 MPSKit = "0.14"

ext/PEPSKitGPUArraysExt.jl

@@ -1,15 +1,137 @@
 module PEPSKitGPUArraysExt
 
 using GPUArrays
-using PEPSKit, MatrixAlgebraKit
+using GPUArrays: AnyGPUArray, AllocCache
+using PEPSKit
 using LinearAlgebra
+using TensorKit
+using TensorKit: MatrixAlgebraKit as MAK
 import PEPSKit: _column_vectors, _materialize_block, _indexable_diagview
 
 PEPSKit._materialize_block(m::Union{Adjoint{T, A}, Transpose{T, A}}) where {T, A <: AnyGPUArray} = copyto!(similar(m, size(m)), m)
 function PEPSKit._column_vectors(::Type{A}, m::AbstractMatrix) where {A <: AnyGPUArray}
     m̃ = PEPSKit._materialize_block(m)
     return A[m̃[:, j] for j in axes(m̃, 2)]
 end
-PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MatrixAlgebraKit.diagview(m))
+PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MAK.diagview(m))
+
+# Each caller (such as `su_iter`) gets a pair of caches. This makes sense to do on a per-caller basis
+# because what is being cached varies between algorithms.
+# For each caller we also store several caches, for SimultaneousCTMRG and SU,
+# one for even iterations and one for odd, and for SequentialCTMRG, 5 (one "round" plus one extra)
+# This has to be done because we can't reuse a cache from iteration `i`
+# until iteration `i+n` is completely finished and its result handed off.
+const ALLOC_CACHES = Dict{Tuple{Symbol, Int}, Vector{AllocCache}}()
+const ALLOC_CACHES_LOCK = ReentrantLock()
+
+function _caches(site::Symbol, depth::Int)
+    return Base.@lock ALLOC_CACHES_LOCK begin
+        get!(() -> [AllocCache() for _ in 1:depth], ALLOC_CACHES, (site, depth))
+    end
+end
+
+function PEPSKit._with_alloc_cache(f, ::Type{<:AnyGPUArray}, site::Symbol, iter::Int, depth::Int)
+    cache = @inbounds _caches(site, depth)[mod1(iter + 1, depth)]
+    return GPUArrays.@cached cache f()
+end
+
+# Reduce into a 0-dimensional device array instead of returning a host scalar. `sdiag_pow` only
+# feeds this into a broadcast, and a 0-dim array broadcasts as a scalar, so the value never has to
+# come back to the host. Returning a number here would force a device sync on every call, and
+# `sdiag_pow` runs once per bond per weight absorption in simple update.
+function PEPSKit._maxabs(data::AnyGPUArray)
+    T = real(eltype(data))
+    acc = similar(data, T, ())
+    fill!(acc, zero(T))
+    Base.mapreducedim!(abs, max, acc, data)
+    return acc
+end
+
+PEPSKit._uncache(x, ::Type{<:AnyGPUArray}) = deepcopy(x)
+
+function PEPSKit.free_alloc_caches!(::Type{<:AnyGPUArray}, caller::Symbol)
+    Base.@lock ALLOC_CACHES_LOCK begin
+        # collect first: freeing mutates ALLOC_CACHES
+        stale = [key for key in keys(ALLOC_CACHES) if first(key) === caller]
+        for key in stale
+            for cache in ALLOC_CACHES[key]
+                GPUArrays.unsafe_free!(cache)
+            end
+            delete!(ALLOC_CACHES, key)
+        end
+    end
+    return nothing
+end
+
+function PEPSKit.free_alloc_caches!(::Type{<:AnyGPUArray})
+    Base.@lock ALLOC_CACHES_LOCK begin
+        for caches in values(ALLOC_CACHES), cache in caches
+            GPUArrays.unsafe_free!(cache)
+        end
+        empty!(ALLOC_CACHES)
+    end
+    return nothing
+end
+
+"""
+    _batched_spectra_alg(proto) -> alg or nothing
+
+Default batched SVD algorithm this backend offers, or `nothing` if it has none.
+"""
+function _batched_spectra_alg(proto)
+    # TODO BAD MAKE THIS A MAK CALL
+    alg = try
+        MAK.default_svd_algorithm(typeof(similar(proto, 0, 0, 0)))
+    catch
+        return nothing
+    end
+    return alg isa MAK.AbstractAlgorithm ? alg : nothing
+end
+
+# Hook into the collection-level convergence API. Deliberately restricted to the generic
+# CTMRG algorithms: `C4vCTMRG` overrides `corner_spectrum` to `eigh_vals` (its corners are
+# diagonal), so a blanket override here would silently switch it back to `svd_vals`.
+function PEPSKit.corner_spectra(
+        Cs::AbstractArray{<:AbstractTensorMap},
+        ::Union{PEPSKit.SequentialCTMRG, PEPSKit.SimultaneousCTMRG},
+    )
+    return _batched_spectra(Cs)
+end
+function PEPSKit.edge_spectra(
+        Ts::AbstractArray{<:AbstractTensorMap},
+        ::Union{PEPSKit.SequentialCTMRG, PEPSKit.SimultaneousCTMRG},
+    )
+    return _batched_spectra(Ts)
+end
+
+# `calc_convergence` decomposes every corner and every edge of the environment
+# for regular CTMRG, which is expensive, at *least* 8 separate `svd_vals` calls.
+# Across *multiple tensors* the situation is much better than within *one*,
+# because the corners have to all share a space,
+# so for a given sector their blocks have identical sizes and batch with no padding at all.
+function _batched_spectra(ts::AbstractArray{T}) where {T <: AbstractTensorMap}
+    # TODO BAD FIND A BETTER DISPATCH HERE
+    (isempty(ts) || !(TensorKit.storagetype(T) <: AnyGPUArray)) && return map(svd_vals, ts)
+    items = [(i, c) for i in eachindex(ts) for c in blocksectors(ts[i])]
+    isempty(items) && return map(svd_vals, ts)
+    alg = _batched_spectra_alg(block(ts[first(items)[1]], first(items)[2]))
+    isnothing(alg) && return map(svd_vals, ts)
+
+    Ss = map(
+        t -> MAK.initialize_output(
+            MAK.svd_vals!, t,
+            MAK.default_algorithm(
+                MAK.svd_vals!, typeof(t)
+            )
+        ), ts
+    )
+    # `batched_svd_vals!` destroys the blocks it has to decompose one at a time, but `ts` is the
+    # live environment, and computing the convergence spectra must not damage the environment
+    # it is measuring. Copying whole tensors costs one copy per tensor instead of one per block.
+    ts′ = map(copy, ts)
+    As = [block(ts′[i], c) for (i, c) in items]
+    MAK.batched_svd_vals!(As, [block(Ss[i], c) for (i, c) in items], alg)
+    return Ss
+end
 
 end

src/PEPSKit.jl

@@ -23,7 +23,7 @@ using KrylovKit
 using KrylovKit: Lanczos, BlockLanczos
 
 using TensorOperations, OptimKit
-using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add
+using TensorOperations: AbstractBackend, DefaultBackend, DefaultAllocator, tensoralloc_add
 using TensorKit.Strided: StridedView, sreshape
 import OptimKit: fixedpoint
 using ChainRulesCore, Zygote
@@ -60,6 +60,7 @@ include("Defaults.jl")  # Include first to allow for docstring interpolation wit
 include("utility/util.jl")
 include("utility/contraction_labels.jl")
 include("utility/tensor_traces.jl")
+include("utility/alloc_cache.jl")
 include("utility/indexing.jl")
 include("utility/diffable_threads.jl")
 include("utility/twistdual.jl")

src/algorithms/ctmrg/c4v.jl

@@ -248,7 +248,7 @@ function initialize_singlet_c4v_env(state::InfinitePEPS, Venv::ElementarySpace)
 end
 function initialize_singlet_c4v_env(T, state::InfinitePEPS, Venv::ElementarySpace)
     Vpeps = north_virtualspace(state, 1, 1)'
-    return initialize_singlet_c4v_env(similarstoragetype(storagetype(state), eltype(T)), Vpeps, Venv)
+    return initialize_singlet_c4v_env(similarstoragetype(storagetype(state), real(eltype(T))), Vpeps, Venv)
 end
 function initialize_singlet_c4v_env(T, Vpeps::ElementarySpace, Venv::ElementarySpace)
     realT = similarstoragetype(similarstoragetype(T), real(eltype(T)))

src/algorithms/ctmrg/ctmrg.jl

@@ -49,6 +49,12 @@ Perform a single CTMRG iteration in which all directions are being grown and ren
 """
 function ctmrg_iteration(network, env, alg::CTMRGAlgorithm) end
 
+# Signature of the buffer sizes a CTMRG run allocates: the corner and edge spaces fix them,
+# and the edges carry the network bond dimension as well as the environment dimension.
+function _ctmrg_cache_signature(env::CTMRGEnv)
+    return hash((alloc_cache_signature(env.corners), alloc_cache_signature(env.edges)))
+end
+
 """
     leading_boundary(env₀, network; kwargs...) -> env, info
     # expert version:
@@ -109,6 +115,11 @@ function leading_boundary(
         env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm
     )
     check_input(leading_boundary, network, env₀, alg)
+    # cached buffer sizes are set by the corner/edge spaces, and the edges carry the network
+    # bond dimension too, so a change here means every pooled buffer has gone stale
+    ignore_derivatives() do
+        free_stale_alloc_caches!(storagetype(env₀), :ctmrg, :enter, _ctmrg_cache_signature(env₀))
+    end
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
     return with_verbosity(; alg.verbosity) do
         env = deepcopy(env₀)
@@ -120,7 +131,10 @@ function leading_boundary(
         local info_iter
         converged = false
         for iter in 1:(alg.maxiter)
-            env, info_iter = ctmrg_iteration(network, env, alg)
+            env, info_iter = with_alloc_cache(storagetype(env), :ctmrg, iter, alloc_cache_depth(alg)) do
+                ctmrg_iteration(network, env, alg)
+            end
+            env = uncache(env, storagetype(env))
             η, CS, TS = calc_convergence(env, CS, TS, alg)
 
             if η ≤ alg.tol && iter ≥ alg.miniter
@@ -134,6 +148,15 @@ function leading_boundary(
                 ctmrg_logiter!(log, iter, η, network, env)
             end
         end
+        env = uncache(env, storagetype(env))
+        # a truncation that is not fixed-space grows the corner/edge spaces as the loop runs,
+        # leaving one pooled buffer set per intermediate shape. `env` is copied out by now, so
+        # nothing live is backed by the cache and it is safe to drop those here. Re-checking
+        # the signature keeps the pool warm for the repeated fixed-space calls of an
+        # optimization loop, where the spaces do not move.
+        ignore_derivatives() do
+            free_stale_alloc_caches!(storagetype(env), :ctmrg, :exit, _ctmrg_cache_signature(env))
+        end
         info = (;
             converged,
             convergence_error = η,
@@ -242,6 +265,7 @@ edge_spectra(Ts, alg) = map(T -> edge_spectrum(T, alg), Ts)
 The spectrum of a corner, used to measure CTMRG convergence.
 """
 corner_spectrum(C, alg) = svd_vals(C)
+
 """
     edge_spectrum(T, alg)
 

src/algorithms/time_evolution/simpleupdate.jl

@@ -162,7 +162,11 @@ end
 function Base.iterate(it::TimeEvolver{<:SimpleUpdate}, state = it.state)
     iter, t = state.iter, state.t
     (iter == it.nstep) && return nothing
-    psi, env, ϵ = su_iter(state.psi, it.circuit, it.alg, state.env)
+    storage = storagetype(state.psi)
+    psi, env, ϵ = with_alloc_cache(storage, :su, iter) do
+        su_iter(state.psi, it.circuit, it.alg, state.env)
+    end
+    psi, env = uncache(psi, storage), uncache(env, storage)
     # update internal state
     iter += 1
     t += it.dt