From 0c2ec9f4ef83a6d6ae9dc4dad767b18e6a35755e Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Mon, 2 Nov 2026 20:03:08 +0800 Subject: [PATCH 1/2] Use allocators for multi-site gates in simple update --- src/PEPSKit.jl | 2 +- src/algorithms/time_evolution/apply_mpo.jl | 64 +++++++++++-------- src/algorithms/time_evolution/simpleupdate.jl | 2 +- .../time_evolution/simpleupdate3site.jl | 36 +++++++++-- 4 files changed, 72 insertions(+), 32 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index e3b683987..6668964aa 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -23,7 +23,7 @@ using KrylovKit using KrylovKit: Lanczos, BlockLanczos using TensorOperations, OptimKit -using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add +using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add, tensoralloc_contract using TensorKit.Strided: StridedView, sreshape import OptimKit: fixedpoint using ChainRulesCore, Zygote diff --git a/src/algorithms/time_evolution/apply_mpo.jl b/src/algorithms/time_evolution/apply_mpo.jl index 511efb47d..df614ab44 100644 --- a/src/algorithms/time_evolution/apply_mpo.jl +++ b/src/algorithms/time_evolution/apply_mpo.jl @@ -120,21 +120,26 @@ Perform QR decomposition through a `GenericMPSTensor` ``` """ function qr_through( - R0::MPSBondTensor, M::GenericMPSTensor{S, N}; normalize::Bool = true + R0::MPSBondTensor, M::GenericMPSTensor{S, N}; normalize::Bool = true, allocator = DefaultAllocator() ) where {S, N} @assert !isdual(codomain(R0, 1)) @assert !isdual(domain(M, 1)) && !isdual(codomain(M, 1)) pR = (codomainind(R0), domainind(R0)) pM = ((1,), Tuple(2:(N + 1))) pRM = (codomainind(M), domainind(M)) - A = tensorcontract(R0, pR, false, M, pM, false, pRM) + cp = allocator_checkpoint!(allocator) + TC = promote_type(scalartype(R0), scalartype(M)) + A = tensoralloc_contract(TC, R0, pR, false, M, pM, false, pRM, Val(true), allocator) + tensorcontract!(A, R0, pR, false, M, pM, false, pRM, One(), Zero(), DefaultBackend(), allocator) _, r = left_orth!(A; positive = true) + tensorfree!(A, allocator) + allocator_reset!(allocator, cp) normalize && normalize!(r, Inf) return r end # for `M` at the left end of the MPS function qr_through( - ::Nothing, M::GenericMPSTensor{S, N}; normalize::Bool = true + ::Nothing, M::GenericMPSTensor{S, N}; normalize::Bool = true, allocator = DefaultAllocator() ) where {S, N} @assert !isdual(domain(M, 1)) _, r = left_orth(M; positive = true) @@ -151,21 +156,26 @@ Perform LQ decomposition through a `GenericMPSTensor` ``` """ function lq_through( - M::GenericMPSTensor{S, N}, L1::MPSBondTensor; normalize::Bool = true + M::GenericMPSTensor{S, N}, L1::MPSBondTensor; normalize::Bool = true, allocator = DefaultAllocator() ) where {S, N} @assert !isdual(domain(L1, 1)) @assert !isdual(codomain(M, 1)) && !isdual(domain(M, 1)) pM = (codomainind(M), domainind(M)) pL = (codomainind(L1), domainind(L1)) pML = ((1,), ntuple(i -> i + 1, N)) - A = tensorcontract(M, pM, false, L1, pL, false, pML) + cp = allocator_checkpoint!(allocator) + TC = promote_type(scalartype(M), scalartype(L1)) + A = tensoralloc_contract(TC, M, pM, false, L1, pL, false, pML, Val(true), allocator) + tensorcontract!(A, M, pM, false, L1, pL, false, pML, One(), Zero(), DefaultBackend(), allocator) l, _ = right_orth!(A; positive = true) + tensorfree!(A, allocator) + allocator_reset!(allocator, cp) normalize && normalize!(l, Inf) return l end # for `M` at the right end of the MPS function lq_through( - M::GenericMPSTensor{S, N}, ::Nothing; normalize::Bool = true + M::GenericMPSTensor{S, N}, ::Nothing; normalize::Bool = true, allocator = DefaultAllocator() ) where {S, N} @assert !isdual(codomain(M, 1)) A = permute(M, ((1,), ntuple(i -> i + 1, N)); copy = true) @@ -177,20 +187,20 @@ end """ Given a cluster `Ms`, find all `R`, `L` matrices on each internal bond """ -function _get_allRLs(Ms::Vector{T}) where {T <: GenericMPSTensor} +function _get_allRLs(Ms::Vector{T}; allocator = DefaultAllocator()) where {T <: GenericMPSTensor} # M1 -- (R1,L1) -- M2 -- (R2,L2) -- M3 N = length(Ms) # get the first R and the last L - R_first = qr_through(nothing, first(Ms); normalize = true) - L_last = lq_through(last(Ms), nothing; normalize = true) + R_first = qr_through(nothing, first(Ms); normalize = true, allocator) + L_last = lq_through(last(Ms), nothing; normalize = true, allocator) Rs = Vector{typeof(R_first)}(undef, N - 1) Ls = Vector{typeof(L_last)}(undef, N - 1) Rs[1], Ls[end] = R_first, L_last # get remaining R, L matrices for n in 2:(N - 1) m = N - n + 1 - Rs[n] = qr_through(Rs[n - 1], Ms[n]; normalize = true) - Ls[m - 1] = lq_through(Ms[m], Ls[m]; normalize = true) + Rs[n] = qr_through(Rs[n - 1], Ms[n]; normalize = true, allocator) + Ls[m - 1] = lq_through(Ms[m], Ls[m]; normalize = true, allocator) end return Rs, Ls end @@ -224,10 +234,10 @@ Given a cluster `Ms`, find all projectors `Pa`, `Pb` and Schmidt weights `wts` on internal bonds. """ function _get_allprojs( - Ms::Vector{T}, truncs::Vector{E} + Ms::Vector{T}, truncs::Vector{E}; allocator = DefaultAllocator() ) where {T <: GenericMPSTensor, E <: TruncationStrategy} N = length(Ms) - Rs, Ls = _get_allRLs(Ms) + Rs, Ls = _get_allRLs(Ms; allocator) @assert length(truncs) == N - 1 projs_errs = map(Rs, Ls, truncs) do R, L, trunc return _proj_from_RL(R, L; trunc) @@ -273,9 +283,9 @@ end Find projectors to truncate internal bonds of the cluster `Ms`. """ function _cluster_truncate!( - Ms::Vector{T}, truncs::Vector{E} + Ms::Vector{T}, truncs::Vector{E}; allocator = DefaultAllocator() ) where {T <: GenericMPSTensor, E <: TruncationStrategy} - Pas, Pbs, wts, ϵs = _get_allprojs(Ms, truncs) + Pas, Pbs, wts, ϵs = _get_allprojs(Ms, truncs; allocator) # apply projectors # M1 -- (Pa1,wt1,Pb1) -- M2 -- (Pa2,wt2,Pb2) -- M3 for (i, (Pa, Pb)) in enumerate(zip(Pas, Pbs)) @@ -300,7 +310,8 @@ e.g. Cluster in PEPS with `gate_ax = 1`: ``` """ function _apply_gatempo!( - Ms::Vector{T1}, gs::Vector{T2}; gate_ax::Int = 1 + Ms::Vector{T1}, gs::Vector{T2}; + gate_ax::Int = 1, allocator = DefaultAllocator() ) where {T1 <: GenericMPSTensor{<:ElementarySpace, 4}, T2 <: AbstractTensorMap} @assert length(Ms) == length(gs) @assert gate_ax == 1 @@ -324,20 +335,21 @@ function _apply_gatempo!( @assert !isdual(space(M, 2)) if i == 1 fr = fusers[i] - @tensor (Ms[i])[-1 -2 -3 -4; -5] := M[-1 1 -3 -4; 2] * g[-2 1 3] * fr'[2 3; -5] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4; -5] := M[-1 1 -3 -4; 2] * g[-2 1 3] * fr'[2 3; -5] elseif i == length(Ms) fl = fusers[i - 1] - @tensor (Ms[i])[-1 -2 -3 -4; -5] := fl[-1; 2 3] * M[2 1 -3 -4; -5] * g[3 -2 1] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4; -5] := fl[-1; 2 3] * M[2 1 -3 -4; -5] * g[3 -2 1] else fl, fr = fusers[i - 1], fusers[i] - @tensor (Ms[i])[-1 -2 -3 -4; -5] := fl[-1; 2 3] * M[2 1 -3 -4; 4] * g[3 -2 1 5] * fr'[4 5; -5] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4; -5] := fl[-1; 2 3] * M[2 1 -3 -4; 4] * g[3 -2 1 5] * fr'[4 5; -5] end end return Ms end function _apply_gatempo!( - Ms::Vector{T1}, gs::Vector{T2}; gate_ax::Int = 1 + Ms::Vector{T1}, gs::Vector{T2}; + gate_ax::Int = 1, allocator = DefaultAllocator() ) where {T1 <: GenericMPSTensor{<:ElementarySpace, 5}, T2 <: AbstractTensorMap} @assert length(Ms) == length(gs) @assert gate_ax == 1 || gate_ax == 2 @@ -373,23 +385,23 @@ function _apply_gatempo!( if i == 1 fr = fusers[i] if gate_ax == 1 - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := M[-1 1 -3 -4 -5; 2] * g[-2 1 3] * fr'[2 3; -6] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := M[-1 1 -3 -4 -5; 2] * g[-2 1 3] * fr'[2 3; -6] else - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := M[-1 -2 1 -4 -5; 2] * g[1 -3 3] * fr'[2 3; -6] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := M[-1 -2 1 -4 -5; 2] * g[1 -3 3] * fr'[2 3; -6] end elseif i == length(Ms) fl = fusers[i - 1] if gate_ax == 1 - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 1 -3 -4 -5; -6] * g[3 -2 1] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 1 -3 -4 -5; -6] * g[3 -2 1] else - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 -2 1 -4 -5; -6] * g[3 1 -3] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 -2 1 -4 -5; -6] * g[3 1 -3] end else fl, fr = fusers[i - 1], fusers[i] if gate_ax == 1 - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 1 -3 -4 -5; 4] * g[3 -2 1 5] * fr'[4 5; -6] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 1 -3 -4 -5; 4] * g[3 -2 1 5] * fr'[4 5; -6] else - @tensor (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 -2 1 -4 -5; 4] * g[3 1 -3 5] * fr'[4 5; -6] + @tensor allocator = allocator (Ms[i])[-1 -2 -3 -4 -5; -6] := fl[-1; 2 3] * M[2 -2 1 -4 -5; 4] * g[3 1 -3 5] * fr'[4 5; -6] end end end diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 38a3614ea..ff5cd133c 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -152,7 +152,7 @@ function su_iter( else # N-site MPO gate (N ≥ 2) alg.bipartite && error("Multi-site MPO gates are not compatible with bipartite states.") - ϵ′ = _su_iter!(state2, gate, env2, sites, alg) + ϵ′ = _su_iter!(state2, gate, env2, sites, alg; allocator) ϵ = max(ϵ, ϵ′) end end diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl index 1614f6e79..c2f5da3c3 100644 --- a/src/algorithms/time_evolution/simpleupdate3site.jl +++ b/src/algorithms/time_evolution/simpleupdate3site.jl @@ -29,11 +29,37 @@ function _get_cluster_with_weights( return Ms_mps, open_vaxs, invperms end +""" +Factor spectator legs out of the cluster endpoints, retaining unit legs to preserve the MPS tensor layout. +""" +function _reduce_mpo_endpoints!(Ms::Vector{<:GenericMPSTensor}) + N = numout(first(Ms)) + physical = ntuple(i -> i + 1, N - 3) + Ql, R = left_orth!(permute(first(Ms), ((1, N - 1, N), (physical..., N + 1))); positive = true) + R = permute(R, ((1, physical...), (N - 1,))) + L, Qr = right_orth!(permute(last(Ms), ((1, physical...), (N - 1, N, N + 1))); positive = true) + Ms[1] = insertrightunit(insertrightunit(R, Val(N - 2)), Val(N - 1)) + Ms[end] = insertrightunit(insertrightunit(L, Val(N - 2)), Val(N - 1)) + return Ql, Qr +end + +""" +Restore the spectator legs of the reduced cluster endpoints after MPO application and truncation. +""" +function _restore_mpo_endpoints!(Ms::Vector{<:GenericMPSTensor}, Ql, Qr) + N = numout(first(Ms)) + R = removeunit(removeunit(first(Ms), Val(N)), Val(N - 1)) + L = removeunit(removeunit(last(Ms), Val(N)), Val(N - 1)) + p = ((1, ntuple(i -> i + 3, N - 3)..., 2, 3), (N + 1,)) + Ms[1] = tensorcontract(Ql, ((1, 2, 3), (4,)), false, R, ((1,), Tuple(2:(N - 1))), false, p) + Ms[end] = tensorcontract(L, (codomainind(L), domainind(L)), false, Qr, (codomainind(Qr), domainind(Qr)), false, (ntuple(identity, N), (N + 1,))) + return Ms +end + """ Simple update with an N-site MPO `gate` (N ≥ 2). -The `allocator` is accepted for compatibility with the nearest neighbor update, -but is not used yet. +Temporary MPO contractions and QR/LQ inputs use `allocator`. """ function _su_iter!( state::InfiniteState, gate::Vector{T}, env::SUWeight, @@ -45,15 +71,17 @@ function _su_iter!( # flip virtual arrows in `Ms` to ← flips = [isdual(space(M, 1)) for M in Iterators.drop(Ms, 1)] _flip_virtuals!(Ms, flips) + Ql, Qr = _reduce_mpo_endpoints!(Ms) # apply gate MPOs and truncate ϵ = 0.0 local wts for gate_ax in 1:2 - _apply_gatempo!(Ms, gate; gate_ax) - wts, ϵs, = _cluster_truncate!(Ms, truncs) + _apply_gatempo!(Ms, gate; gate_ax, allocator) + wts, ϵs, = _cluster_truncate!(Ms, truncs; allocator) ϵ = max(ϵ, maximum(ϵs)) alg.purified && break end + _restore_mpo_endpoints!(Ms, Ql, Qr) # restore virtual arrows in `Ms` _flip_virtuals!(Ms, flips) # update env weights From 11b27b85d1b8083a4cf74f07197315a333ef5616 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Mon, 2 Nov 2026 20:06:56 +0800 Subject: [PATCH 2/2] Docstring rewording --- src/algorithms/time_evolution/simpleupdate3site.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl index c2f5da3c3..82b3d1118 100644 --- a/src/algorithms/time_evolution/simpleupdate3site.jl +++ b/src/algorithms/time_evolution/simpleupdate3site.jl @@ -30,7 +30,8 @@ function _get_cluster_with_weights( end """ -Factor spectator legs out of the cluster endpoints, retaining unit legs to preserve the MPS tensor layout. +Factor out irrelevant virtual legs out of the cluster endpoints +for more efficient application of gate MPOs. """ function _reduce_mpo_endpoints!(Ms::Vector{<:GenericMPSTensor}) N = numout(first(Ms)) @@ -44,7 +45,7 @@ function _reduce_mpo_endpoints!(Ms::Vector{<:GenericMPSTensor}) end """ -Restore the spectator legs of the reduced cluster endpoints after MPO application and truncation. +Restore virtual legs of the cluster endpoints after gate MPO application. """ function _restore_mpo_endpoints!(Ms::Vector{<:GenericMPSTensor}, Ql, Qr) N = numout(first(Ms))