Avoid copies in the c4v sign flip calculation

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},

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

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