From a2d071227932c45676ee12f5ace3385dcc363cc8 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Wed, 4 Nov 2026 22:48:49 +0800 Subject: [PATCH 1/3] TimeEvolver-owned allocator --- src/algorithms/time_evolution/simpleupdate.jl | 20 +++++++++---- src/algorithms/time_evolution/time_evolve.jl | 11 ++++++-- test/testsuite/timeevol/timestep.jl | 28 +++++++++++++++++++ 3 files changed, 52 insertions(+), 7 deletions(-) diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index ff5cd133c..81b992afe 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -38,6 +38,10 @@ struct SUState{S <: InfiniteState, E <: SUWeight, N <: Number} env::E end +# Also covers trivial simple update used by SUGauge. +TimeEvolver(alg::SimpleUpdate, dt::Number, nstep::Int, circuit, state::SUState) = + TimeEvolver(alg, dt, nstep, circuit, state, default_allocator(state.psi, SerialScheduler())) + """ TimeEvolver( psi0::Union{InfinitePEPS, InfinitePEPO}, H::LocalOperator, dt::Number, @@ -118,15 +122,16 @@ function _su_iter!( end """ -One iteration of simple update +One iteration of simple update. Pass `allocator` to reuse contraction scratch; +`TimeEvolver` supplies its owned allocator automatically. """ function su_iter( state::InfiniteState, circuit::LocalCircuit, - alg::SimpleUpdate, env::SUWeight + alg::SimpleUpdate, env::SUWeight; + allocator = default_allocator(state, SerialScheduler()) ) state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0 - # the gates are applied one after another, so a single allocator serves all of them - allocator = default_allocator(state, SerialScheduler()) + # Gates are applied sequentially and share the same scratch allocator. for (sites, gate) in circuit.gates if length(sites) == 1 # 1-site gate @@ -162,7 +167,12 @@ end function Base.iterate(it::TimeEvolver{<:SimpleUpdate}, state = it.state) iter, t = state.iter, state.t (iter == it.nstep) && return nothing - psi, env, ϵ = su_iter(state.psi, it.circuit, it.alg, state.env) + cp = allocator_checkpoint!(it.allocator) + psi, env, ϵ = try + su_iter(state.psi, it.circuit, it.alg, state.env; allocator = it.allocator) + finally + allocator_reset!(it.allocator, cp) + end # update internal state iter += 1 t += it.dt diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl index d6065d20a..c55a04a88 100644 --- a/src/algorithms/time_evolution/time_evolve.jl +++ b/src/algorithms/time_evolution/time_evolve.jl @@ -6,7 +6,7 @@ Abstract super type for time evolution algorithms of InfinitePEPS or InfinitePEP abstract type TimeEvolution end """ - mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number} + mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number, A} Iterator for Trotter-based time evolution of InfinitePEPS or InfinitePEPO. @@ -14,7 +14,7 @@ Iterator for Trotter-based time evolution of InfinitePEPS or InfinitePEPO. $(TYPEDFIELDS) """ -mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number} +mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number, A} "Time evolution algorithm (currently supported: `SimpleUpdate`)" alg::TE "Trotter time step" @@ -26,8 +26,15 @@ mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number} "Internal state of the iterator, including the number of already performed iterations, evolved time, PEPS/PEPO and its environment" state::S + "Reusable contraction scratch retained across timesteps and state resets. + An evolver must not be advanced concurrently from multiple tasks." + allocator::A end +# Preserve construction for algorithms that do not yet reuse contraction scratch. +TimeEvolver(alg::TimeEvolution, dt::Number, nstep::Int, circuit, state) = + TimeEvolver(alg, dt, nstep, circuit, state, DefaultAllocator()) + Base.iterate(it::TimeEvolver) = iterate(it, it.state) """ diff --git a/test/testsuite/timeevol/timestep.jl b/test/testsuite/timeevol/timestep.jl index 7a125b435..2675eb74b 100644 --- a/test/testsuite/timeevol/timestep.jl +++ b/test/testsuite/timeevol/timestep.jl @@ -15,6 +15,7 @@ function timeevol_timestep(AT) dt, nstep = 1.0e-2, 50 # manual timestep evolver = TimeEvolver(ψ0, H, dt, nstep, alg, env0) + allocator = evolver.allocator ψ1, env1, info1 = deepcopy(ψ0), deepcopy(env0), nothing for iter in 0:(nstep - 1) ψ1, env1, info1 = timestep(evolver, ψ1, env1) @@ -31,5 +32,32 @@ function timeevol_timestep(AT) @test ψ1 == ψ2 == ψ3 @test env1 == env2 == env3 @test info1 == info2 == info3 + @test evolver.allocator === allocator + timeevol_su_allocator(AT) + end +end + +"Check fresh/reused scratch equivalence, growth, output ownership, and restarts." +function timeevol_su_allocator(AT) + H = adapt(AT, real(heisenberg_XYZ(ComplexF64, Trivial, InfiniteSquare(2, 2)))) + return @testset "Persistent scratch ($State, MPO=$force_mpo, $AT)" for + State in (InfinitePEPS, InfinitePEPO), force_mpo in (false, true) + ψ0 = adapt(AT, State(randn, Float64, fill(ℂ^2, 2, 2), fill(ℂ^2, 2, 2))) + env0 = SUWeight(ψ0) + alg = SimpleUpdate(; trunc = truncrank(4), force_mpo, purified = (State === InfinitePEPS)) + it = TimeEvolver(ψ0, H, 0.01, 4, alg, env0) + allocator = it.allocator + ψ, env, info = saved = timestep(it, ψ0, env0) + @test (ψ, env, info.ϵ) == PEPSKit.su_iter(ψ0, it.circuit, it.alg, env0) + snapshot = deepcopy(saved) + for _ in 2:4 + expected = PEPSKit.su_iter(ψ, it.circuit, it.alg, env) + ψ, env, info = timestep(it, ψ, env) + @test (ψ, env, info.ϵ) == expected + @test saved == snapshot # later scratch reuse must not overwrite earlier outputs + end + ψ, env, info = timestep(it, ψ0, env0; iter = 0, t = 0.0) + @test (ψ, env, info.ϵ) == PEPSKit.su_iter(ψ0, it.circuit, it.alg, env0) + @test it.allocator === allocator end end From 0e3cf0e97874fcd54da5dfce2559ce4c8507fc1b Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 06:24:47 +0800 Subject: [PATCH 2/3] Remove redundant allocator checkpoint from SU iteration --- src/algorithms/time_evolution/simpleupdate.jl | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 81b992afe..0e18e43b4 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -167,12 +167,7 @@ end function Base.iterate(it::TimeEvolver{<:SimpleUpdate}, state = it.state) iter, t = state.iter, state.t (iter == it.nstep) && return nothing - cp = allocator_checkpoint!(it.allocator) - psi, env, ϵ = try - su_iter(state.psi, it.circuit, it.alg, state.env; allocator = it.allocator) - finally - allocator_reset!(it.allocator, cp) - end + psi, env, ϵ = su_iter(state.psi, it.circuit, it.alg, state.env; allocator = it.allocator) # update internal state iter += 1 t += it.dt From f54f25854f61954cbc7850053e795019f8f70d86 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 06:44:05 +0800 Subject: [PATCH 3/3] Test deduplication --- test/testsuite/timeevol/timestep.jl | 30 ++++------------------------- 1 file changed, 4 insertions(+), 26 deletions(-) diff --git a/test/testsuite/timeevol/timestep.jl b/test/testsuite/timeevol/timestep.jl index 2675eb74b..f765e8562 100644 --- a/test/testsuite/timeevol/timestep.jl +++ b/test/testsuite/timeevol/timestep.jl @@ -25,6 +25,9 @@ function timeevol_timestep(AT) # for-loop syntax # Restart the exhausted evolver using the public timestep keywords. ψ3, env3, info3 = timestep(evolver, ψ0, env0; iter = 0, t = 0.0) + saved = (ψ3, env3, info3) + snapshot = deepcopy(saved) + @test (ψ3, env3, info3.ϵ) == PEPSKit.su_iter(ψ0, evolver.circuit, evolver.alg, env0) for state in evolver ψ3, env3, info3 = state end @@ -33,31 +36,6 @@ function timeevol_timestep(AT) @test env1 == env2 == env3 @test info1 == info2 == info3 @test evolver.allocator === allocator - timeevol_su_allocator(AT) - end -end - -"Check fresh/reused scratch equivalence, growth, output ownership, and restarts." -function timeevol_su_allocator(AT) - H = adapt(AT, real(heisenberg_XYZ(ComplexF64, Trivial, InfiniteSquare(2, 2)))) - return @testset "Persistent scratch ($State, MPO=$force_mpo, $AT)" for - State in (InfinitePEPS, InfinitePEPO), force_mpo in (false, true) - ψ0 = adapt(AT, State(randn, Float64, fill(ℂ^2, 2, 2), fill(ℂ^2, 2, 2))) - env0 = SUWeight(ψ0) - alg = SimpleUpdate(; trunc = truncrank(4), force_mpo, purified = (State === InfinitePEPS)) - it = TimeEvolver(ψ0, H, 0.01, 4, alg, env0) - allocator = it.allocator - ψ, env, info = saved = timestep(it, ψ0, env0) - @test (ψ, env, info.ϵ) == PEPSKit.su_iter(ψ0, it.circuit, it.alg, env0) - snapshot = deepcopy(saved) - for _ in 2:4 - expected = PEPSKit.su_iter(ψ, it.circuit, it.alg, env) - ψ, env, info = timestep(it, ψ, env) - @test (ψ, env, info.ϵ) == expected - @test saved == snapshot # later scratch reuse must not overwrite earlier outputs - end - ψ, env, info = timestep(it, ψ0, env0; iter = 0, t = 0.0) - @test (ψ, env, info.ϵ) == PEPSKit.su_iter(ψ0, it.circuit, it.alg, env0) - @test it.allocator === allocator + @test saved == snapshot # later scratch reuse must not overwrite earlier outputs end end