Refactor simple update and reduce weight-absorption allocations

benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl

@@ -14,9 +14,11 @@ const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json")
 const SUITE = BenchmarkGroup()
 const MODULES = Dict{String, Symbol}(
     "ctmrg" => :CTMRGBenchmarks,
+    "timeevol" => :TimeEvolBenchmarks,
 )
 
 include("ctmrg/CTMRGBenchmarks.jl")
+include("timeevol/TimeEvolBenchmarks.jl")
 
 load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...)
 

benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl

@@ -0,0 +1,76 @@
+module TimeEvolBenchmarks
+
+using BenchmarkTools
+using TensorKit
+using PEPSKit
+using PEPSKit.MPSKitModels: hubbard_space
+
+const SUITE = BenchmarkGroup()
+
+"""
+A simple update scenario on a 2 x 2 unit cell: the Hamiltonian `H` and the initial state `psi`,
+evolved with bond dimension `D` (the truncation keeps the initial virtual spaces).
+"""
+struct SUSpec{H, S}
+    H::H
+    psi::S
+    D::Int
+    purified::Bool
+end
+
+function heisenberg_peps(D::Int; T = Float64)
+    H = heisenberg_XYZ(InfiniteSquare(2, 2))
+    psi = InfinitePEPS(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2))
+    return SUSpec(H, psi, D, true)
+end
+
+function heisenberg_u1_peps(n::Int; T = Float64)
+    H = heisenberg_XXZ(T, U1Irrep, InfiniteSquare(2, 2))
+    V = U1Space(0 => 2n, 1 // 2 => n, -1 // 2 => n)
+    psi = InfinitePEPS(randn, T, physicalspace(H)[1, 1], V; unitcell = (2, 2))
+    return SUSpec(H, psi, dim(V), true)
+end
+
+function hubbard_peps(n::Int; T = Float64)
+    H = hubbard_model(T, Trivial, U1Irrep, InfiniteSquare(2, 2); t = 1.0, U = 6.0, mu = 3.0)
+    V = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2n, (1, 1 // 2) => n, (1, -1 // 2) => n)
+    psi = InfinitePEPS(randn, T, hubbard_space(Trivial, U1Irrep), V; unitcell = (2, 2))
+    return SUSpec(H, psi, dim(V), true)
+end
+
+function heisenberg_pepo(D::Int; T = Float64)
+    H = heisenberg_XYZ(InfiniteSquare(2, 2))
+    psi = InfinitePEPO(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2, 1))
+    return SUSpec(H, psi, D, false)
+end
+
+"""
+Benchmark one simple update iteration, i.e. the application of all Trotter gates of `spec.H`,
+starting from weights obtained from a few warm-up iterations.
+"""
+function su_iter_benchmark(spec::SUSpec; dt = 1.0e-2, nwarmup = 3)
+    alg = SimpleUpdate(; trunc = truncrank(spec.D), spec.purified)
+    circuit = PEPSKit.trotterize(spec.H, dt)
+    psi, env = spec.psi, SUWeight(spec.psi)
+    for _ in 1:nwarmup
+        psi, env, = PEPSKit.su_iter(psi, circuit, alg, env)
+    end
+    return @benchmarkable PEPSKit.su_iter($psi, $circuit, $alg, $env)
+end
+
+const SCENARIOS = [
+    ("Heisenberg PEPS", "Trivial", "D8") => () -> heisenberg_peps(8),
+    ("Heisenberg PEPS", "Trivial", "D16") => () -> heisenberg_peps(16),
+    ("Heisenberg PEPS", "U1", "D8") => () -> heisenberg_u1_peps(2),
+    ("Heisenberg PEPS", "U1", "D16") => () -> heisenberg_u1_peps(4),
+    ("Hubbard PEPS", "fZ2xU1", "D8") => () -> hubbard_peps(2),
+    ("Hubbard PEPS", "fZ2xU1", "D16") => () -> hubbard_peps(4),
+    ("Heisenberg PEPO", "Trivial", "D4") => () -> heisenberg_pepo(4),
+    ("Heisenberg PEPO", "Trivial", "D6") => () -> heisenberg_pepo(6),
+]
+
+for (groupkeys, spec) in SCENARIOS
+    SUITE[("su_iter", groupkeys...)] = su_iter_benchmark(spec())
+end
+
+end # module

benchmark/README.md

@@ -1,9 +1,13 @@
 # PEPSKit benchmarks
 
 A [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl)-based suite that exercises
-performance-critical paths in PEPSKit. The suite is organised into modules; the only module
-currently registered is `ctmrg`, which times a single `ctmrg_iteration` across several CTMRG
-algorithms, problem scenarios, and symmetry sectors.
+performance-critical paths in PEPSKit. The suite is organised into modules:
+
+- `ctmrg` times a single `ctmrg_iteration` across several CTMRG algorithms, problem scenarios,
+  and symmetry sectors.
+- `timeevol` times a single simple update iteration (`su_iter`, i.e. one sweep of all Trotter
+  gates over a 2 x 2 unit cell) for PEPS and PEPO time evolution, across models, symmetry sectors
+  and bond dimensions.
 
 ## Layout
 
@@ -16,11 +20,13 @@ bench/benchmark/
 └── PEPSKitBenchmarks/
     ├── PEPSKitBenchmarks.jl      # module registry (MODULES), load!, loadall!
     ├── utils/BenchUtils.jl       # TOML <-> VectorSpace conversion helpers
-    └── ctmrg/
-        ├── CTMRGBenchmarks.jl    # populates SUITE["ctmrg"]
-        ├── ctmrg_iteration_benchmarks.jl
-        ├── default.toml          # Trivial + NonUniform scenarios
-        └── su3_hubbard.toml      # SU(3) fermionic scenarios
+    ├── ctmrg/
+    │   ├── CTMRGBenchmarks.jl    # populates SUITE["ctmrg"]
+    │   ├── ctmrg_iteration_benchmarks.jl
+    │   ├── default.toml          # Trivial + NonUniform scenarios
+    │   └── su3_hubbard.toml      # SU(3) fermionic scenarios
+    └── timeevol/
+        └── TimeEvolBenchmarks.jl # populates SUITE["timeevol"]
 ```
 
 New modules are registered in the `MODULES` dict in `PEPSKitBenchmarks.jl`.
@@ -54,6 +60,18 @@ julia --project=bench/benchmark bench/benchmark/benchmarks.jl --modules=ctmrg
 The script only populates `SUITE`; call `run(SUITE)` afterwards (e.g. from an interactive
 session) to execute. Available IDs are the keys of `MODULES` in `PEPSKitBenchmarks.jl`.
 
+### Simple update
+
+`SUITE["timeevol"]` is indexed as `"su_iter" → model → symmetry → "D<D>"`, e.g.
+
+```julia
+include("bench/benchmark/benchmarks.jl")
+run(SUITE["timeevol"]["su_iter"]["Hubbard PEPS"]["fZ2xU1"]["D16"])
+```
+
+The scenarios are defined in the `SCENARIOS` list in `timeevol/TimeEvolBenchmarks.jl`; each
+entry maps its group keys to a function that builds the Hamiltonian and initial state.
+
 ### By group key
 
 `SUITE` is a nested `BenchmarkGroup` indexed as:

src/PEPSKit.jl

@@ -23,6 +23,8 @@ using KrylovKit
 using KrylovKit: Lanczos, BlockLanczos
 
 using TensorOperations, OptimKit
+using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add
+using TensorKit.Strided: StridedView, sreshape
 import OptimKit: fixedpoint
 using ChainRulesCore, Zygote
 import TupleTools
@@ -32,6 +34,7 @@ using MPSKit: MPSTensor, MPOTensor, GenericMPSTensor, MPSBondTensor, ProductTran
 using MPSKit: InfiniteEnvironments
 using MPSKit: DynamicTol
 using MPSKit.DynamicTols: adapt_solver
+using MPSKit: default_allocator
 using MPSKit: @log_initialization, @log_iteration, @log_convergence, @log_nonconvergence
 using MPSKit.IterativeLoggers: with_verbosity
 import MPSKit.DynamicTols: _updatetol
@@ -49,7 +52,7 @@ import MPSKitModels:
     transverse_field_ising, heisenberg_XYZ, heisenberg_XXZ,
     hubbard_model, bose_hubbard_model, tj_model
 using FiniteDifferences
-using OhMyThreads: tmap, tmap!
+using OhMyThreads: tmap, tmap!, SerialScheduler
 using DocStringExtensions
 
 include("Defaults.jl")  # Include first to allow for docstring interpolation with Defaults values

src/algorithms/contractions/absorb_weight.jl

@@ -48,18 +48,58 @@ function absorb_weight(
         t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight,
         row::Int, col::Int, virt_axes::NTuple{N, Int}; inv::Bool = false
     ) where {N}
-    Np = numout(t)
-    vax = first(virt_axes)
-    wt = weight_to_absorb(weights, row, col, vax; inv)
-    axes, t2 = absorb_first_weight(t, wt, vax)
-    for vax in Base.tail(virt_axes)
-        axes, biperm = _permute_to_last(axes, vax + Np)
-        wt = weight_to_absorb(weights, row, col, vax; inv)
-        # use `*` to make absorption/removal twist-free
-        t2 = permute(t2, biperm) * wt
+    wts = map(ax -> weight_to_absorb(weights, row, col, ax; inv), virt_axes)
+    T = promote_type(scalartype(t), map(scalartype, wts)...)
+    return _absorb_weights!(similar(t, T), t, wts, virt_axes)
+end
+
+"""
+    absorb_weight!(t, weights::SUWeight, rowcol::CartesianIndex{2}, virt_axes; inv = false)
+
+In-place version of [`absorb_weight`](@ref).
+"""
+function absorb_weight!(
+        t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight,
+        rowcol::CartesianIndex{2}, virt_axes::NTuple{N, Int}; inv::Bool = false
+    ) where {N}
+    wts = map(ax -> weight_to_absorb(weights, rowcol[1], rowcol[2], ax; inv), virt_axes)
+    return _absorb_weights!(t, t, wts, virt_axes)
+end
+
+# Write `src` into `dst` with `weights[i]` absorbed into virtual axis `virtual_axes[i]`,
+# where `domain(weights[i]) == domain(src)[virtual_axes[i]]`; `dst === src` is allowed.
+# A diagonal weight only rescales the basis vectors of its leg, by a factor that depends on
+# the sector of the leg and the index within it. This commutes with any permutation, so it is
+# applied directly on each fusion-tree block of `src` in its own leg order; it is equivalent
+# to `permute(src, ...) * wt` followed by the inverse permutation.
+# Since the domain of `wt` is the space of the leg itself (dual legs get `transpose(wt)` from
+# `weight_to_absorb`), `wt` is composed onto the leg without bending it, so even for dual legs
+# no twist is involved.
+function _absorb_weights!(
+        dst::AbstractTensorMap, src::AbstractTensorMap,
+        weights::NTuple{N, DiagonalTensorMap}, virtual_axes::NTuple{N, Int}
+    ) where {N}
+    Np, M = numout(src), Val(numind(src)) # static rank, so the reshaped weights are inferred
+    space(dst) == space(src) ||
+        throw(SpaceMismatch(lazy"destination $(space(dst)) does not match source $(space(src))"))
+    for (wt, ax) in zip(weights, virtual_axes)
+        domain(wt)[1] == domain(src)[ax] ||
+            throw(SpaceMismatch(lazy"weight $(space(wt)) does not match virtual axis $ax of $(space(src))"))
+    end
+    # Look up every block of the weights only once, reshaped to broadcast along their leg, and
+    # iterate over the subblocks of `dst` and `src` (which have the same order) with `subblocks`,
+    # as the per-block `block(wt, c)` and `t[f₁, f₂]` would both redo the structure lookups.
+    λdicts = map(weights, virtual_axes) do wt, ax
+        return Dict(
+            c => sreshape(StridedView(b.diag), ntuple(i -> i == Np + ax ? length(b.diag) : 1, M))
+                for (c, b) in blocks(wt)
+        )
     end
-    perm_back = invperm(axes)
-    return permute(t2, (perm_back[begin:numout(t)], perm_back[(numout(t) + 1):end]))
+    for (((f₁, f₂), d), (_, s)) in zip(subblocks(dst), subblocks(src))
+        λs = map((λdict, ax) -> λdict[f₂.uncoupled[ax]], λdicts, virtual_axes)
+        @. d = *(s, λs...)
+    end
+    return dst
 end
 
 """
@@ -86,13 +126,3 @@ function weight_to_absorb(
     (ax == SOUTH || ax == WEST) && return transpose(wt)
     return wt
 end
-
-function absorb_first_weight(
-        t::Union{PEPSTensor, PEPOTensor}, wt::DiagonalTensorMap, vax::Int
-    )
-    Np = numout(t)
-    axes, biperm = _permute_to_last(ntuple(identity, numind(t)), vax + Np)
-    # use `*` to make absorption/removal twist-free
-    t2 = permute(t, biperm) * wt
-    return axes, t2
-end

src/algorithms/time_evolution/apply_gate.jl

@@ -20,33 +20,112 @@ function _apply_sitegate(
     return a′
 end
 
+# --- 2-site gates on nearest neighbor bonds ---------------------------------------------
+#
+# The bond between `s₁` and `s₂` is updated in the native PEPS/PEPO leg order, using the
+# bond direction to locate the legs. Each site tensor is split as `T = Q * R`, where the
+# reduced tensor `R` holds the physical leg acted on by the gate and the bond leg. The gate is
+# applied to `R₁ -- R₂` and split again by a truncated SVD, after which `Q` and the new `R`
+# are contracted back into the native leg order.
+#
+# Temporaries allocated with `allocator` never outlive the function that allocates them:
+# each function below frees its own temporaries and resets `allocator` before returning.
+
+"""
+    _orient_bond(state, sites, gate) -> (s₁, k₁), (s₂, k₂), gate
+
+Order the two `sites` of a nearest neighbor bond and the virtual axes `k₁`, `k₂` of their
+bond legs, such that the bond leg of `s₁` is non-dual. Putting `s₁` on the codomain side
+of the SVD then reproduces the original arrow on the bond, so nothing has to be flipped.
+"""
+function _orient_bond(state::InfiniteState, sites::Vector{CartesianIndex{2}}, gate::NNGate)
+    s₁, s₂ = sites
+    k₁ = first(_nn_bonddir(s₁, s₂))
+    k₂ = mod1(k₁ + 2, 4)
+    T₁ = state[s₁]
+    # `space` of a domain leg returns its dual
+    isdual(space(T₁, numout(T₁) + k₁)) && return (s₁, k₁), (s₂, k₂), gate
+    return (s₂, k₂), (s₁, k₁), permute(gate, ((2, 1), (4, 3)))
+end
+
 """
-$(SIGNATURES)
+    _bond_qr(T, k, gate_ax, [allocator]) -> Q, R, p
 
-Apply 2-site `gate` on the reduced bond tensors `a`, `b`
+Split `T = Q * R`, where the reduced tensor `R : x ← p ⊗ v` holds the physical axis `gate_ax`
+and the virtual axis `k` (the bond) of `T`, while `Q` holds all other axes. `p` is the
+permutation that brings the product back to the leg order of `T`. For example, for a
+PEPSTensor `T[p; n e s w]` with the bond `v = e` on its east axis (`k = 2`),
 ```
-    -1← a --- 3 --- b ← -4          -2         -3
-        ↓           ↓               ↓           ↓
-        1           2               |----gate---|
-        ↓           ↓       or      ↓           ↓
-        |----gate---|               1           2
-        ↓           ↓               ↓           ↓
-        -2         -3           -1← a --- 3 --- b ← -4
+        n                       n
+        |                       |
+    w - T - v       =       w - Q ← x ← R - v
+        | ↘                     |        ↘
+        s   p                   s         p
 ```
+For a PEPOTensor, the physical axis that is not `gate_ax` stays in `Q`.
+The permuted copy of `T` that is factorized is a temporary allocated with `allocator`.
 """
-function _apply_gate(a::MPSTensor, b::MPSTensor, gate::NNGate, trunc::TruncationStrategy)
-    V = space(b, 1)
-    need_flip = isdual(V)
-    if isdual(space(a, 2))
-        @tensor a2b2[-1 -2; -3 -4] := gate[1 2; -2 -3] * a[-1 1 3] * b[3 2 -4]
+function _bond_qr(
+        T::Union{PEPSTensor, PEPOTensor}, k::Int, gate_ax::Int, allocator = DefaultAllocator()
+    )
+    Np, M = numout(T), numind(T)
+    r = (gate_ax, Np + k)
+    q = TupleTools.deleteat(ntuple(identity, M), r)
+    cp = allocator_checkpoint!(allocator)
+    T′ = tensoralloc_add(scalartype(T), T, (q, r), false, Val(true), allocator)
+    tensoradd!(T′, T, (q, r), false, One(), Zero(), DefaultBackend(), allocator)
+    Q, R = left_orth!(T′; positive = true)
+    tensorfree!(T′, allocator)
+    allocator_reset!(allocator, cp)
+    return Q, R, _inv_mpo_perm((q..., r...), Val(Np))
+end
+
+"""
+    _undo_bond_qr(Q, X, pX, p, [allocator])
+
+Inverse of [`_bond_qr`](@ref): contract `Q` with `X`, whose legs selected by `pX` are
+ordered as `x ← p ⊗ v`, into the leg order specified by `p`.
+"""
+function _undo_bond_qr(
+        Q::AbstractTensorMap, X::AbstractTensorMap, pX, p, allocator = DefaultAllocator()
+    )
+    N = numind(Q)
+    pQ = (ntuple(identity, N - 1), (N,))
+    cp = allocator_checkpoint!(allocator)
+    T = tensorcontract(Q, pQ, false, X, pX, false, p, One(), DefaultBackend(), allocator)
+    allocator_reset!(allocator, cp)
+    return T
+end
+
+"""
+    _apply_bond_gate(R₁, R₂, gate, trunc, [allocator]) -> U, S, Vᴴ, ϵ
+
+Apply the 2-site `gate` to the reduced tensors `R₁ : x₁ ← p₁ ⊗ v` and `R₂ : x₂ ← p₂ ⊗ v'`
+```
+    x₁ ← R₁ --- v --- R₂ → x₂           p₁            p₂
+         ↘            ↘                 ↓             ↓
+          p₁           p₂               |----gate-----|
+          ↓            ↓        or      ↓             ↓
+          |----gate----|                p₁            p₂
+          ↓            ↓                ↘             ↘
+          p₁           p₂          x₁ ← R₁ --- v --- R₂ → x₂
+```
+where the second form is used when the physical legs are dual (the second physical axis of
+a PEPO), and split the result as `U S Vᴴ` with a truncated SVD. The square root of `S` is
+absorbed into `U : x₁ ⊗ p₁ ← v′` and `Vᴴ : v′ ← x₂ ⊗ p₂`.
+"""
+function _apply_bond_gate(
+        R₁::AbstractTensorMap, R₂::AbstractTensorMap, gate::NNGate,
+        trunc::TruncationStrategy, allocator = DefaultAllocator()
+    )
+    θ = if isdual(space(R₁, 2))
+        @tensor allocator = allocator θ[-1 -2; -3 -4] :=
+            R₁[-1; 1 3] * R₂[-3; 2 3] * gate[1 2; -2 -4]
     else
-        @tensor a2b2[-1 -2; -3 -4] := gate[-2 -3; 1 2] * a[-1 1 3] * b[3 2 -4]
-    end
-    a, s, b, ϵ = svd_trunc!(a2b2; trunc)
-    a, b = absorb_s(a, s, b)
-    if need_flip
-        a, s, b = flip(a, numind(a)), _fliptwist_s!(s), flip(b, 1)
+        @tensor allocator = allocator θ[-1 -2; -3 -4] :=
+            R₁[-1; 1 3] * R₂[-3; 2 3] * gate[-2 -4; 1 2]
     end
-    b = permute(b, ((1, 2), (3,)))
-    return a, s, b, ϵ
+    U, S, Vᴴ, ϵ = svd_trunc!(θ; trunc)
+    absorb_s!(U, S, Vᴴ)
+    return U, S, Vᴴ, ϵ
 end