GPU optimization for maxabs

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

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

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

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 *= _maxabs(s)
     spow = sdiag_pow(s, pow; tol)
     spow_minus1_conj = scale!(sdiag_pow(s', pow - 1; tol), pow)
     function sdiag_pow_pullback(c̄_)