diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 78834a705..5d276c8d9 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -1,6 +1,7 @@ module PEPSKitGPUArraysExt using GPUArrays +using GPUArrays.KernelAbstractions: @kernel, @index, @Const, get_backend using PEPSKit using LinearAlgebra using TensorKit @@ -14,12 +15,44 @@ function PEPSKit._column_vectors(::Type{A}, m::AbstractMatrix) where {A <: AnyGP end PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MAK.diagview(m)) -# 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`. -# Restricted to `BlasFloat` since that is what the batched GPU SVD drivers support. const GPUTensorMap = TensorMap{<:LinearAlgebra.BlasFloat, <:Any, <:Any, <:Any, <:AbstractGPUArray} +_offsets(S::TensorKit.SectorDict) = Int32.(first.(values(S))) +_offsets(S::AbstractDict) = sort!(Int32.(first.(values(S)))) +function PEPSKit.corner_spectrum(C::GPUTensorMap, ::PEPSKit.C4vCTMRG) + S = eigh_vals(C) + x = parent(S) + isempty(x) && return S + m = reduce(_maxabs_signed, x; dims = 1, init = zero(eltype(x))) + starts = _offsets(S.structure) + push!(starts, length(x) + 1) + offsets = copyto!(similar(x, Int32, length(starts)), starts) + _flip_spectrum_kernel!(get_backend(x))(x, offsets, m; ndrange = length(x)) + return S +end +_maxabs_signed(a, b) = abs(a) > abs(b) ? a : abs(b) > abs(a) ? b : max(a, b) + +# if `m[1] < 0`, negate `x` and reverse it within each block; `offsets` are the sorted block start +# positions followed by `length(x) + 1` +@kernel function _flip_spectrum_kernel!(x, @Const(offsets), @Const(m)) + i = @index(Global, Linear) + @inbounds if m[1] < zero(eltype(m)) + lo, hi = 1, length(offsets) - 1 # last block starting at or before `i` + while lo < hi + mid = (lo + hi + 1) >> 1 + offsets[mid] <= i ? (lo = mid) : (hi = mid - 1) + end + j = offsets[lo] + offsets[lo + 1] - 1 - i # mirror image of `i` within its block + if i < j + xi, xj = x[i], x[j] + x[i] = -xj + x[j] = -xi + elseif i == j + x[i] = -x[i] + end + end +end + function PEPSKit.corner_spectra( Cs::AbstractArray{<:GPUTensorMap}, ::Union{PEPSKit.SequentialCTMRG, PEPSKit.SimultaneousCTMRG}, diff --git a/test/ctmrg/c4v_sign_branches.jl b/test/ctmrg/c4v_sign_branches.jl index 10a5dca95..228c2509b 100644 --- a/test/ctmrg/c4v_sign_branches.jl +++ b/test/ctmrg/c4v_sign_branches.jl @@ -1,4 +1,4 @@ -using PEPSKit +using PEPSKit, CUDA, AMDGPU @isdefined(TestSuite) || include("../testsuite/TestSuite.jl") using .TestSuite @@ -8,3 +8,11 @@ is_buildkite = get(ENV, "BUILDKITE", "false") == "true" if !is_buildkite TestSuite.ctmrg_c4v_sign_branches(Vector) end + +if CUDA.functional() + TestSuite.ctmrg_c4v_sign_branches(CuArray) +end + +if AMDGPU.functional() + TestSuite.ctmrg_c4v_sign_branches(ROCArray) +end diff --git a/test/testsuite/ctmrg/c4v_sign_branches.jl b/test/testsuite/ctmrg/c4v_sign_branches.jl index 4b2567464..396c703c2 100644 --- a/test/testsuite/ctmrg/c4v_sign_branches.jl +++ b/test/testsuite/ctmrg/c4v_sign_branches.jl @@ -16,9 +16,9 @@ Venv = U1Space(0 => 2, 1 // 2 => 2, -1 // 2 => 2, 1 => 1, -1 => 1) tol = 1.0e-10 maxiter = 400 -# (-1)^{2Q} on a U(1) space -function charge_parity(W) - X = id(W) +# (-1)^{2Q} on a U(1) space, with storage type `A` +function charge_parity(A, W) + X = id(A, W) for (c, b) in blocks(X) b .*= (-1)^Int(2 * c.charge) end @@ -35,7 +35,7 @@ end function sign_twin(env) E = env.edges[1] - g = charge_parity(space(E, 2)) + g = charge_parity(storagetype(E), space(E, 2)) @tensor Ê[χ_out D_ket D_bra; χ_in] := E[χ_out D_ket′ D_bra; χ_in] * g[D_ket; D_ket′] return C4vCTMRGEnv(env.corners[1], Ê) end @@ -56,6 +56,7 @@ function ctmrg_c4v_sign_branches(AT) @test iteration_sign(network, env_twin, alg) == -iteration_sign(network, env, alg) env_twin, info_twin = leading_boundary(env_twin, peps, alg) @test info_twin.converged - @test expectation_value(peps, H, env_twin) ≈ expectation_value(peps, H, env) atol = 1.0e-8 + H_AT = adapt(AT, H) + @test expectation_value(peps, H_AT, env_twin) ≈ expectation_value(peps, H_AT, env) atol = 1.0e-8 end end