[WIP] GPU support in PEPSKit

.buildkite/pipeline.yml

@@ -1,59 +1,37 @@
 env:
-  SECRET_CODECOV_TOKEN: "MH6hHjQi7vG2V1Yfotv5/z5Dkx1k5SdyGYlGTFXiQr22XksJgsXaBuvFKUrjC7JwcpBsOVU8103LuMKl3m7VJ35WzHZrOssYycVbdGcb2kloc6xvUOsN2R5BrhCQ4Pii0l6ZeVRjCnZVkcmb0Rf4glGFyfibCrqniry8RLhblsuFKFsijRK4OxiWYEs1IvUulN+ER8tEsEtw4+ZqC5nbLGMSnUG/saPkDQOVIBscvikbKEnBcCXBheGPktF+Y/cy/1Xa+FiBPoZcypwTeAjKG1g0MqyHXjaYekb/7fekaj+hukGaeJSCXxY8KEb2IZCh+Y36Tp6y6qsIp/AdtEnCpQ==;U2FsdGVkX18WQxvGLspPwzC4aDe+U7TXU+itebTbgh8LUkE6GukxxReHYiDZ6IrBiVvSGTVJMquW0c8KsOI1pw=="
+  SECRET_CODECOV_TOKEN: "yXWVfuSYtwx4ksmqquaMEj0TWi42b1QZxPEwJccOBQ0fSZqNdkqGS/Z2fZ4bEPOoUGEzFPlqn75ZVe9nwf4XlbZCN0mPyYcA3DQiqAwO+9SskfwolCYBfI2RsBsVlZy7YXDHRH5KTGvelU0/fuuCt/DAk2j+Xk9HOZr4kx5RxEG1dKvBzUGB8q5phgJjvm0DUQ3w12iJMVeQVWU02P6dmrDPdiUJzoTzF0Gqpo4ZaKBfK9u58WP5Ao6vwNcffVoOiy4fiH93oXExOKoc2dqlKeAyLBHVXAy6wtKxZvNOrigyqKp0UEdmhkIe4iNNzo3bg3o0AIaS3HbkvPaU0ttg7Q==;U2FsdGVkX18oOnMGdD2pf5EYcmGkA11S5cIoWyXhsCDO9HMj9r0sNxou8l5EkreJ2FZDJs1AaIMf1X9tHs3WnQ=="
 
 steps:
-  - label: "Julia v1 -- CUDA"
+  - label: "Julia {{matrix.julia}} -- {{matrix.queue}} / {{matrix.group}}"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1"
-      - JuliaCI/julia-test#v1: ~
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "{{matrix.group}}"
       - JuliaCI/julia-coverage#v1:
           dirs:
             - src
             - ext
     agents:
-      queue: "cuda"
+      queue: "{{matrix.queue}}"
     if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
+    timeout_in_minutes: 120
+    matrix:
+      setup:
+        julia:
+          - "1.10"
+          - "1.13"
+        queue:
+          - "cuda"
+          - "rocm"
+        group:
+          - "boundarymps"
+          - "bondenv"
+          - "bp"
+          - "compress"
+          - "ctmrg"
+          - "gradients"
+          - "timeevol"
+          - "toolbox"
+          - "utility"
 
-  - label: "Julia LTS -- CUDA"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1.10" # "lts" isn't valid
-      - JuliaCI/julia-test#v1: ~
-      - JuliaCI/julia-coverage#v1:
-          dirs:
-            - src
-            - ext
-    agents:
-      queue: "cuda"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
-
-  - label: "Julia v1 -- AMDGPU"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1"
-      - JuliaCI/julia-test#v1: ~
-      - JuliaCI/julia-coverage#v1:
-          dirs:
-            - src
-            - ext
-    agents:
-      queue: "rocm"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
-
-  - label: "Julia LTS -- AMDGPU"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1.10" # "lts" isn't valid
-      - JuliaCI/julia-test#v1: ~
-      - JuliaCI/julia-coverage#v1:
-          dirs:
-            - src
-            - ext
-    agents:
-      queue: "rocm"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60

Project.toml

@@ -23,15 +23,32 @@ OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
+StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorKitTensors = "41b62e7d-e9d1-4e23-942c-79a97adf954b"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+
+[sources]
+MPSKit = {rev = "main", url = "https://github.com/QuantumKitHub/MPSKit.jl"}
+MatrixAlgebraKit = {rev = "ksh/gesvdx-rank1-guard", url = "https://github.com/QuantumKitHub/MatrixAlgebraKit.jl"}
+TensorKit = {rev = "ksh/batched_svd", url = "https://github.com/QuantumKitHub/TensorKit.jl"}
+
+[extensions]
+PEPSKitAdaptExt = "Adapt"
+PEPSKitGPUArraysExt = "GPUArrays"
+
 [compat]
 Accessors = "0.1"
+Adapt = "4"
+GPUArrays = "11"
 ChainRulesCore = "1.0"
 Compat = "3.46, 4.2"
 DocStringExtensions = "0.9.3"
@@ -49,8 +66,8 @@ Random = "1"
 Statistics = "1"
 TensorKit = "0.16.5, 0.17"
 TensorKitTensors = "0.3.1"
-TensorOperations = "5"
+TensorOperations = "5.8.1"
 TupleTools = "1.6.0"
-VectorInterface = "0.4, 0.5, 0.6"
+VectorInterface = "0.4, 0.5, 0.6, 0.7"
 Zygote = "0.6, 0.7"
 julia = "1.10"

ext/PEPSKitAdaptExt.jl

@@ -0,0 +1,32 @@
+module PEPSKitAdaptExt
+
+using PEPSKit
+using Adapt
+
+function Adapt.adapt_structure(to, x::PEPSKit.LocalOperator{T, S}) where {T, S}
+    terms′ = Dict(k => adapt(to, v) for (k, v) in x.terms)
+    return PEPSKit.LocalOperator{valtype(terms′)}(x.lattice, terms′)
+end
+
+function Adapt.adapt_structure(to, x::PEPSKit.InfinitePEPS{T}) where {T}
+    A′ = map(a -> adapt(to, a), x.A)
+    return InfinitePEPS{eltype(A′)}(A′)
+end
+
+function Adapt.adapt_structure(to, x::PEPSKit.InfinitePEPO{T}) where {T}
+    A′ = map(a -> adapt(to, a), x.A)
+    return InfinitePEPO{eltype(A′)}(A′)
+end
+
+function Adapt.adapt_structure(to, x::PEPSKit.InfinitePartitionFunction{T}) where {T}
+    A′ = map(a -> adapt(to, a), x.A)
+    return InfinitePartitionFunction{eltype(A′)}(A′)
+end
+
+function Adapt.adapt_structure(to, x::PEPSKit.CTMRGEnv{C, T}) where {C, T}
+    C′ = map(c -> adapt(to, c), x.corners)
+    T′ = map(t -> adapt(to, t), x.edges)
+    return CTMRGEnv{eltype(C′), eltype(T′)}(C′, T′)
+end
+
+end

ext/PEPSKitGPUArraysExt.jl

@@ -0,0 +1,179 @@
+module PEPSKitGPUArraysExt
+
+using GPUArrays
+using GPUArrays: AnyGPUArray, AllocCache
+using PEPSKit
+using TensorKit
+using TensorKit: MatrixAlgebraKit as MAK
+
+# 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 truncated SVD of a whole cluster's internal bonds to avoid multiple small kernel launches.
+function PEPSKit.bond_svds(
+        ::Type{<:AnyGPUArray}, rls::AbstractVector, truncs::AbstractVector
+    )
+    isempty(rls) && return map(_ -> nothing, rls)
+    # The different GPU libaries offer different batching algos,
+    # make sure we have one that actually works.
+    alg = MAK.default_algorithm(MAK.batched_svd_compact!, eltype(rls))
+    Fs = map(rl -> MAK.initialize_output(MAK.svd_compact!, rl, alg), rls)
+    balg = _cluster_batched_alg(rls)
+    if isnothing(balg)
+        for (rl, F) in zip(rls, Fs)
+            MAK.svd_compact!(rl, F, alg)
+        end
+    else
+        # Pool every (bond, sector) block into one ragged batch. MatrixAlgebraKit batches
+        # blocks of equal size together even across different bonds, since the decomposition
+        # does not care which bond a block came from, and zero-pads the leftovers.
+        items = [(i, c) for i in eachindex(rls) for c in blocksectors(rls[i])]
+        As = [block(rls[i], c) for (i, c) in items]
+        Us = [block(Fs[i][1], c) for (i, c) in items]
+        Ss = [TensorKit.diagview(block(Fs[i][2], c)) for (i, c) in items]
+        Vᴴs = [block(Fs[i][3], c) for (i, c) in items]
+        MAK.batched_svd_compact!(As, (Us, Ss, Vᴴs), balg)
+    end
+    return map(Fs, truncs) do F, trunc
+        (U, S, Vᴴ) = F
+        USVᴴtrunc, ind = MAK.truncate(MAK.svd_trunc!, (U, S, Vᴴ), trunc)
+        ϵ = MAK.truncation_error!(TensorKit.diagview(S), ind)
+        return (USVᴴtrunc..., ϵ)
+    end
+end
+
+"""
+    CLUSTER_BATCHED_SVD[]
+
+Whether simple update batches the SVDs of a cluster's internal bonds into one call.
+Off by default.
+"""
+const CLUSTER_BATCHED_SVD = Ref(false)
+
+# Which batched algorithm the backend offers for the cluster's blocks, or `nothing`.
+function _cluster_batched_alg(rls::AbstractVector)
+    CLUSTER_BATCHED_SVD[] || return nothing
+    for i in eachindex(rls), c in blocksectors(rls[i])
+        return _batched_spectra_alg(block(rls[i], c))
+    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/Defaults.jl

@@ -52,7 +52,7 @@ Module containing default algorithm parameter values and arguments.
 
 ## `eigh` forward & reverse
 
-* `eigh_fwd_alg=:$(Defaults.eigh_fwd_alg)` : `eigh` algorithm that is used in the forward pass.
+* `eigh_fwd_alg=:$(Defaults.eigh_fwd_alg)` : `eigh` algorithm that is used in the forward pass. **Note** that on GPU, `DivideAndConquer` is much more performant than `QRIteration`.
     - `:DefaultAlgorithm` : MatrixAlgebraKit's default Eigh algorithm for a given matrix type.
     - `:DivideAndConquer` : MatrixAlgebraKit's [`DivideAndConquer`](@extref MatrixAlgebraKit.DivideAndConquer)
     - `:QRIteration` : MatrixAlgebraKit's [`QRIteration`](@extref MatrixAlgebraKit.QRIteration)

src/PEPSKit.jl

@@ -18,11 +18,14 @@ using TensorKit
 using TensorKit: AdjointTensorMap, SectorDict
 using TensorKit: throw_invalid_innerproduct, similarstoragetype
 using TensorKit.Factorizations: TruncationSpace, _notrunc_ind
+import TensorKit: storagetype
 
 using KrylovKit
 using KrylovKit: Lanczos, BlockLanczos
 
-using TensorOperations, OptimKit
+using TensorOperations
+using TensorOperations: AbstractBackend, DefaultBackend, DefaultAllocator
+using OptimKit
 using ChainRulesCore, Zygote
 using LoggingExtras
 import TupleTools
@@ -54,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/bp/beliefpropagation.jl

@@ -39,6 +39,7 @@ function leading_boundary(env₀::BPEnv, network::InfiniteSquareNetwork, alg::Be
         ϵ = Inf
         @infov 1 loginit!(log, ϵ)
         for iter in 1:(alg.maxiter)
+            # TODO investigate why caching doesn't help here and actually makes things worse
             env′ = bp_iteration(network, env, alg)
             ϵ = oftype(ϵ, tr_distance(env, env′))
             env = env′

src/algorithms/bp/gaugefix.jl

@@ -129,7 +129,7 @@ function SUWeight(env::BPEnv)
         I = CartesianIndex(mod1(dir′ + 1, 2), row, col)
         sqrtM12, _, sqrtM21, _ = _sqrt_bp_messages(I, env)
         Λ = DiagonalTensorMap(svd_vals!(sqrtM12 * sqrtM21))
-        return isdual(space(sqrtM12, 1)) ? _fliptwist_s(Λ) : Λ
+        return isdual(space(sqrtM12, 1)) ? _fliptwist_s!(Λ) : Λ
     end
     return SUWeight(wts)
 end