From 2639908ec813e44bb842ba57e420b61d0a701935 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 6 Nov 2026 05:46:22 -0500 Subject: [PATCH 1/5] Avoid copies in the c4v sign flip calculation --- src/algorithms/ctmrg/c4v.jl | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/algorithms/ctmrg/c4v.jl b/src/algorithms/ctmrg/c4v.jl index cc10c6216..0d4a995ff 100644 --- a/src/algorithms/ctmrg/c4v.jl +++ b/src/algorithms/ctmrg/c4v.jl @@ -74,14 +74,26 @@ end # gauge fix sign of spectrum such that the eigenvalue of largest magnitude is positive function corner_spectrum(C::AbstractTensorMap, ::C4vCTMRG) S = eigh_vals(C) - if -minimum(S) > maximum(S) - for (_, b) in blocks(S) - b .= .-reverse(b) # spectrum of -C: negated and, to stay ascending, reversed - end + flip = _spectrum_sign_flip(S) + for (_, b) in blocks(S) + b .= ifelse.(flip, .-reverse(b), b) end return S end +# Compute whether `-minimum(S) > maximum(S)`, but keep the result as a 1-element array +# with the same storage type as `S`. Avoids synchronizing copies. +function _spectrum_sign_flip(S) + lo = hi = nothing + for (_, b) in blocks(S) + isempty(b) && continue + l, h = view(b, firstindex(b):firstindex(b)), view(b, lastindex(b):lastindex(b)) + lo = isnothing(lo) ? copy(l) : min.(lo, l) + hi = isnothing(hi) ? copy(h) : max.(hi, h) + end + return .-lo .> hi +end + """ convergence_tensors(env::C4vCTMRGEnv, alg::C4vCTMRG) -> (corners, edges) From 3dafb2115965f39c49d1a642fe8327b3c9e9276f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 6 Nov 2026 08:02:02 -0500 Subject: [PATCH 2/5] Do the spectrum flip in one go --- ext/PEPSKitGPUArraysExt.jl | 38 ++++++++++++++++++++--- src/algorithms/ctmrg/c4v.jl | 20 +++--------- test/ctmrg/c4v_sign_branches.jl | 10 +++++- test/testsuite/ctmrg/c4v_sign_branches.jl | 11 ++++--- 4 files changed, 53 insertions(+), 26 deletions(-) diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 78834a705..754a3ef60 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,11 +15,40 @@ 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} +function PEPSKit.corner_spectrum(C::GPUTensorMap, ::PEPSKit.C4vCTMRG) + S = eigh_vals(C) + x = parent(S) + isempty(x) && return S + m = mapreduce(identity, _maxabs_signed, x; dims = 1, init = zero(eltype(x))) + starts = sort!(Int32[first(r) for r in values(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}, diff --git a/src/algorithms/ctmrg/c4v.jl b/src/algorithms/ctmrg/c4v.jl index 0d4a995ff..cc10c6216 100644 --- a/src/algorithms/ctmrg/c4v.jl +++ b/src/algorithms/ctmrg/c4v.jl @@ -74,26 +74,14 @@ end # gauge fix sign of spectrum such that the eigenvalue of largest magnitude is positive function corner_spectrum(C::AbstractTensorMap, ::C4vCTMRG) S = eigh_vals(C) - flip = _spectrum_sign_flip(S) - for (_, b) in blocks(S) - b .= ifelse.(flip, .-reverse(b), b) + if -minimum(S) > maximum(S) + for (_, b) in blocks(S) + b .= .-reverse(b) # spectrum of -C: negated and, to stay ascending, reversed + end end return S end -# Compute whether `-minimum(S) > maximum(S)`, but keep the result as a 1-element array -# with the same storage type as `S`. Avoids synchronizing copies. -function _spectrum_sign_flip(S) - lo = hi = nothing - for (_, b) in blocks(S) - isempty(b) && continue - l, h = view(b, firstindex(b):firstindex(b)), view(b, lastindex(b):lastindex(b)) - lo = isnothing(lo) ? copy(l) : min.(lo, l) - hi = isnothing(hi) ? copy(h) : max.(hi, h) - end - return .-lo .> hi -end - """ convergence_tensors(env::C4vCTMRGEnv, alg::C4vCTMRG) -> (corners, edges) 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 From 96fa4e82d5bef341d0ff280d5283055b836ba2b6 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 6 Nov 2026 15:09:32 +0100 Subject: [PATCH 3/5] Update ext/PEPSKitGPUArraysExt.jl Co-authored-by: Lukas Devos --- ext/PEPSKitGPUArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 754a3ef60..64fd3ac87 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -20,7 +20,7 @@ function PEPSKit.corner_spectrum(C::GPUTensorMap, ::PEPSKit.C4vCTMRG) S = eigh_vals(C) x = parent(S) isempty(x) && return S - m = mapreduce(identity, _maxabs_signed, x; dims = 1, init = zero(eltype(x))) + m = reduce(_maxabs_signed, x; dims = 1, init = zero(eltype(x))) starts = sort!(Int32[first(r) for r in values(S.structure)]) push!(starts, length(x) + 1) offsets = copyto!(similar(x, Int32, length(starts)), starts) From d131283eb20df3d18d488eef3dba746356135f4e Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sun, 8 Nov 2026 01:58:16 -0500 Subject: [PATCH 4/5] Don't sort if we don't need to --- ext/PEPSKitGPUArraysExt.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 64fd3ac87..41db5007b 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -16,12 +16,15 @@ end PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MAK.diagview(m)) const GPUTensorMap = TensorMap{<:LinearAlgebra.BlasFloat, <:Any, <:Any, <:Any, <:AbstractGPUArray} + +_offsets(S::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 = sort!(Int32[first(r) for r in values(S.structure)]) + 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)) From 7484a419b231aeb4fdf93927c2db5819d1f29ef4 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sun, 8 Nov 2026 14:58:59 +0100 Subject: [PATCH 5/5] Update PEPSKitGPUArraysExt.jl Missing module qualifier --- ext/PEPSKitGPUArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 41db5007b..5d276c8d9 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -17,7 +17,7 @@ PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MAK.diagview(m)) const GPUTensorMap = TensorMap{<:LinearAlgebra.BlasFloat, <:Any, <:Any, <:Any, <:AbstractGPUArray} -_offsets(S::SectorDict) = Int32.(first.(values(S))) +_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)