From 53ba7c04573949834bfc63041375f13c76c77483 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 05:58:36 -0400 Subject: [PATCH 1/6] GPU optimization for maxabs --- ext/PEPSKitGPUArraysExt.jl | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ext/PEPSKitGPUArraysExt.jl b/ext/PEPSKitGPUArraysExt.jl index 13510fd5c..e76b0a3f2 100644 --- a/ext/PEPSKitGPUArraysExt.jl +++ b/ext/PEPSKitGPUArraysExt.jl @@ -12,4 +12,16 @@ function PEPSKit._column_vectors(::Type{A}, m::AbstractMatrix) where {A <: AnyGP end PEPSKit._indexable_diagview(m::AnyGPUArray) = collect(MatrixAlgebraKit.diagview(m)) +# 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 CPU. Returning a scalar 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 + end From 81d68de382d90d8806f572fbaa30becc58bf1a9a Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 29 Oct 2026 15:32:53 -0400 Subject: [PATCH 2/6] Fix maxabs use in rrule --- src/utility/util.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/utility/util.jl b/src/utility/util.jl index df25d2ae5..c3bf03a25 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -37,7 +37,7 @@ function sdiag_pow(s::DiagonalTensorMap, pow::Real; tol::Real = eps(real(scalart return spow end function sdiag_pow( - s::AbstractTensorMap{T, S, 1, 1}, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4) + s::AbstractTensorMap{T, S, 1, 1}, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4) ) where {T, S} tol *= _maxabs(s) spow = similar(s) @@ -53,7 +53,7 @@ function ChainRulesCore.rrule( ::typeof(sdiag_pow), s::AbstractTensorMap, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4), ) - tol *= _maxabs(s) + tol *= norm(s, Inf) spow = sdiag_pow(s, pow; tol) spow_minus1_conj = scale!(sdiag_pow(s', pow - 1; tol), pow) function sdiag_pow_pullback(c̄_) From b8752d318abcc6092a65e286f9a6017bc1f7389f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 3 Nov 2026 04:21:52 -0500 Subject: [PATCH 3/6] Go back to norm(Inf) for trotter --- src/algorithms/time_evolution/trotter_gate.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl index 37997e6b4..9503ca4d2 100644 --- a/src/algorithms/time_evolution/trotter_gate.jl +++ b/src/algorithms/time_evolution/trotter_gate.jl @@ -22,7 +22,7 @@ function gate_to_mpo( # convert to Vidal gauge _cluster_truncate!(Os, fill(notrunc(), N - 1)) # evenly distribute the (Inf) norm - nrms = _maxabs.(Os) + nrms = norm.(Os, Inf) fac = prod(nrms)^(1 / N) scale!.(Os, fac ./ nrms) # remove trivial legs in first/last tensor, and restore MPO convention From 1e46fee531e8d9002281680c9b08d474dd9d7477 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 3 Nov 2026 04:24:45 -0500 Subject: [PATCH 4/6] Formatter --- src/utility/util.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/utility/util.jl b/src/utility/util.jl index c3bf03a25..0c9ae97e9 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -37,7 +37,7 @@ function sdiag_pow(s::DiagonalTensorMap, pow::Real; tol::Real = eps(real(scalart return spow end function sdiag_pow( - s::AbstractTensorMap{T, S, 1, 1}, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4) + s::AbstractTensorMap{T, S, 1, 1}, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4) ) where {T, S} tol *= _maxabs(s) spow = similar(s) From 1e5517df0b990a63e161f6d5f3465b017bea06bc Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 5 Nov 2026 16:26:27 +0100 Subject: [PATCH 5/6] Remove stray line --- src/utility/util.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/utility/util.jl b/src/utility/util.jl index 0c9ae97e9..663c4ae50 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -53,7 +53,6 @@ function ChainRulesCore.rrule( ::typeof(sdiag_pow), s::AbstractTensorMap, pow::Real; tol::Real = eps(real(scalartype(s)))^(3 / 4), ) - tol *= norm(s, Inf) spow = sdiag_pow(s, pow; tol) spow_minus1_conj = scale!(sdiag_pow(s', pow - 1; tol), pow) function sdiag_pow_pullback(c̄_) From 9439ae18011672d96a473757bbf35f212f8454ee Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 6 Nov 2026 01:24:47 -0500 Subject: [PATCH 6/6] Drop notrunc --- src/PEPSKit.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 4280c30af..a8039169a 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -17,7 +17,7 @@ using MatrixAlgebraKit: svd_pullback!, svd_trunc_pullback! using TensorKit using TensorKit: AdjointTensorMap, SectorDict using TensorKit: throw_invalid_innerproduct, similarstoragetype -using TensorKit.Factorizations: TruncationSpace, _notrunc_ind +using TensorKit.Factorizations: TruncationSpace using KrylovKit using KrylovKit: Lanczos, BlockLanczos