diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index e6cbf14eb..04b807804 100644 --- a/.buildkite/pipeline.yml +++ b/.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" diff --git a/Project.toml b/Project.toml index a718d73fd..a1559b05f 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,9 @@ 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" @@ -40,11 +43,11 @@ 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" diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 13510fd5c..608033885 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -1,8 +1,11 @@ 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) @@ -10,6 +13,125 @@ function PEPSKit._column_vectors(::Type{A}, m::AbstractMatrix) where {A <: AnyGP 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 diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 4280c30af..e90b022fb 100644 --- a/src/PEPSKit.jl +++ b/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") diff --git a/src/algorithms/ctmrg/c4v.jl b/src/algorithms/ctmrg/c4v.jl index a16d80f01..12486145c 100644 --- a/src/algorithms/ctmrg/c4v.jl +++ b/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))) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 9263eb89f..9c61f47c1 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/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) diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 38a3614ea..d1a00b1f4 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/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 diff --git a/src/utility/alloc_cache.jl b/src/utility/alloc_cache.jl new file mode 100644 index 000000000..d17b45abe --- /dev/null +++ b/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)) diff --git a/test/testsuite/ctmrg/fixed_iterscheme.jl b/test/testsuite/ctmrg/fixed_iterscheme.jl index 835b8cfb5..37d112731 100644 --- a/test/testsuite/ctmrg/fixed_iterscheme.jl +++ b/test/testsuite/ctmrg/fixed_iterscheme.jl @@ -21,7 +21,7 @@ D = 2 χ = 16 projector_algs_asymm = [:HalfInfiniteProjector] #, :FullInfiniteProjector] unitcells = [(1, 1), (3, 4)] -atol = 1.0e-5 +const atol = 1.0e-5 function ctmrg_fixed_iterscheme_asymmetric(AT; svd_alg = :DivideAndConquer) svd_algs = [(; alg = svd_alg), (; alg = :GKL)]