Use caching for SU and CTMRG algos on GPU

ext/PEPSKitGPUArraysExt.jl

@@ -1,15 +1,63 @@
 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 can also store several caches. For CTMRG and SU,
+# we currently store one but we retain the ability to store a "stack" if needed in the future.
+# TODO could this be better handled by LRU.jl or MemoizationKit.jl?
+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
+
+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
 
 end

src/PEPSKit.jl

@@ -57,6 +57,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/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::AbstractCTMRGEnv)
+    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₀::CTMRGEnv, 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

@@ -158,7 +158,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

src/utility/alloc_cache.jl

@@ -0,0 +1,131 @@
+"""
+    with_alloc_cache(f, storage, caller::Symbol, iter::Int) -> f()
+
+Run one iteration `f()` of an iterative algorithm using a memory cache.
+
+**This is a no-op unless `GPUArrays` is loaded _and_ `storage` is a GPU array.**
+
+There seems to be no benefit to caching for CPU-side memory, but for the device,
+using a warm memory pool which recycles memory avoids everything being blocked
+while the GPU device driver allocates.
+
+`caller` ids the calling algo (`:su`, `:ctmrg`) so that algorithms allocating different
+buffer sizes avoid stepping on each others' caches. Not every algorithm needs an allocation cache,
+only the ones with repeated iterations.
+
+Belief propagation seems to not benefit from the caching as much, so it's currently unused there.
+
+`iter` selects between `depth` alternating caches per `caller`. A buffer allocated during
+iteration `i` can't become reusable until later iterations have completely used and discarded
+its output, so with too few caches it would be possible to overwrite state that is still in
+use. A caller that copies its result out with [`uncache`](@ref) before the next iteration
+begins needs only a single cache; simple update keeps 2, since it hands its result off after
+the following step.
+
+!!! note
+    Buffers are only returned to the cache when the enclosing block exits, so a single
+    block retains *everything* it allocated rather than just its maximum live block.
+    Keeping more caches in the rotation multiplies this effect, which matters for
+    algorithms like (esp. sequential) CTMRG whose iterations allocate many
+    short-lived temporaries.
+
+Caching is skipped while `Zygote.jl` is differentiating, because the reverse-mode tape holds references
+to intermediates, and recycling those could silently corrupt gradients.
+"""
+function with_alloc_cache(f, storage::Type, caller::Symbol, iter::Int, depth::Int = 2)
+    Zygote.isderiving() && return f()
+    return _with_alloc_cache(f, storage, caller, iter, depth)
+end
+
+"""
+    uncache(x, storage::Type) -> x
+
+Copy `x` out of any allocation-cache region, so that it stays live and is not overwriten
+once the enclosing [`with_alloc_cache`](@ref) block has exited. This is necessary to
+ensure `leading_boundary` and other functions which call back into AD handle cached
+memory correctly.
+
+Buffers allocated inside a cache block are handed back to the pool when the block exits,
+and the next cached call may hand them out again, which *silently overwrites* a result the
+caller is still holding. Anything that escapes such a block must be copied out.
+
+**This is a no-op unless `GPUArrays` is loaded _and_ `storage` is a GPU array**, and also
+while Zygote is differentiating, since caching is skipped in both of those cases.
+"""
+function uncache(x, storage::Type)
+    Zygote.isderiving() && return x
+    return _uncache(x, storage)
+end
+_uncache(x, ::Type) = x
+
+"""
+    alloc_cache_depth(alg)
+
+How many iterations a buffer must go unused for before it may be recycled.
+"""
+alloc_cache_depth(alg) = 1
+_with_alloc_cache(f, ::Type, ::Symbol, ::Int, ::Int) = f()
+
+"""
+    free_alloc_caches!(storage)
+    free_alloc_caches!(storage, caller::Symbol)
+
+Release memory held by the allocation caches for storage type `storage`, either for every
+`caller` (if none is provided) or only for the given one. Does nothing unless the
+`PEPSKitGPUArraysExt` extension is loaded and `storage` is a GPU array type.
+This should be called when the bond dimension changes, because the cache keys depend on
+buffer size and the caches can't be reused when the bond dimension has changed.
+
+!!! warning
+    This releases the underlying device memory, so it is only safe to call when nothing
+    still references a buffer that was allocated inside a [`with_alloc_cache`](@ref) block.
+    Results that escape such a block must have been copied out with [`uncache`](@ref)
+    first, otherwise freeing the cache leaves them undefined.
+"""
+free_alloc_caches!(::Type) = nothing
+free_alloc_caches!(::Type, ::Symbol) = nothing
+
+# last buffer-shape signature seen per caller, used to detect when cached buffer sizes
+# have gone stale because a bond dimension changed
+const ALLOC_CACHE_SIGNATURES = Dict{Tuple{Symbol, Symbol}, UInt}()
+const ALLOC_CACHE_SIGNATURES_LOCK = ReentrantLock()
+
+"""
+    free_stale_alloc_caches!(storage, caller::Symbol, phase::Symbol, signature::UInt)
+
+Release `caller`'s allocation caches when `signature` differs from the one seen at the same
+`phase` of the previous call, and record `signature` as the current one for that phase.
+
+The caches are keyed by buffer size, so that when the size changes the old, unusable caches
+can be freed and new ones allocated, corresponding to the new size. This avoids the cache
+size growing without bound.
+
+`phase` distinguishes the points a caller checks from. For example, CTMRG grows the
+corner and edge spaces as it converges, so its incoming and outgoing environments differ
+whenever the truncation is *not* fixed-space. Recording both under one key makes the stored
+signature alternate between them, so every check reports stale and the pool is freed on every
+call. Comparing each phase only against itself keeps the pool across repeated calls while
+still invalidating it when the spaces genuinely change.
+
+Like [`free_alloc_caches!`](@ref) this is a no-op without a GPU storage type, and it's
+skipped while `Zygote.jl` is differentiating, since caching is disabled there anyway.
+"""
+function free_stale_alloc_caches!(storage::Type, caller::Symbol, phase::Symbol, signature::UInt)
+    Zygote.isderiving() && return nothing
+    stale = Base.@lock ALLOC_CACHE_SIGNATURES_LOCK begin
+        key = (caller, phase)
+        previous = get(ALLOC_CACHE_SIGNATURES, key, nothing)
+        ALLOC_CACHE_SIGNATURES[key] = signature
+        !isnothing(previous) && previous != signature
+    end
+    stale && free_alloc_caches!(storage, caller)
+    return nothing
+end
+
+"""
+    alloc_cache_signature(tensors) -> UInt
+
+Hash the spaces of `tensors`, for use as a [`free_stale_alloc_caches!`](@ref) signature.
+Two states whose tensors live in the same spaces allocate the same buffer sizes.
+"""
+alloc_cache_signature(tensors) = hash(map(space, tensors))