More allocator threading

src/algorithms/derivatives/hamiltonian_derivatives.jl

@@ -47,8 +47,8 @@ function AC_hamiltonian(
         backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
     )
     @assert below === above "JordanMPO assumptions break"
-    GL = leftenv(envs, site, below)
-    GR = rightenv(envs, site, below)
+    GL = leftenv(envs, site, below; backend, allocator)
+    GR = rightenv(envs, site, below; backend, allocator)
     W = operator[site]
     H_AC = JordanMPO_AC_Hamiltonian(GL, W, GR; backend, allocator)
     return prepare ? prepare_operator!!(H_AC) : H_AC
@@ -209,8 +209,8 @@ function AC2_hamiltonian(
         backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
     )
     @assert below === above "JordanMPO assumptions break"
-    GL = leftenv(envs, site, below)
-    GR = rightenv(envs, site + 1, below)
+    GL = leftenv(envs, site, below; backend, allocator)
+    GR = rightenv(envs, site + 1, below; backend, allocator)
     W1, W2 = operator[site], operator[site + 1]
     H_AC2 = JordanMPO_AC2_Hamiltonian(GL, W1, W2, GR; backend, allocator)
     return prepare ? prepare_operator!!(H_AC2) : H_AC2

src/algorithms/derivatives/mpo_derivatives.jl

@@ -220,13 +220,85 @@ function prepare_operator!!(H::MPO_C_Hamiltonian{<:MPSTensor, <:MPSTensor})
     rightenv = H.rightenv isa TensorMap ? H.rightenv : TensorMap(H.rightenv)
     return prepared_operator_type(typeof(H))(leftenv, rightenv, H.backend, H.allocator)
 end
-function prepare_operator!!(H::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor})
-    backend, allocator = H.backend, H.allocator
+
+# Scratch space of `prepare_operator!!`
+# ------------------------------------
+# The environments of a prepared derivative are the dense, fused form of an environment-operator
+# contraction. Neither that contraction nor its dense copy is kept - only the repartitioned
+# result is - so both are taken from the allocator and handed straight back.
+#
+# `GL_O` and `O_GR` are allocated here rather than by `:=`, which always allocates the result of
+# a `@plansor` block on the heap.
+
+# A `TensorMap` is already dense, so `repartition` is the only tensor that has to be kept.
+@inline function _fuse_env(t::TensorMap, N₁::Int, N₂::Int, backend, allocator)
+    return repartition(fuse_legs(t, N₁, N₂), 2, 2; copy = true, backend, allocator)
+end
+@inline function _fuse_env(t::AbstractBlockTensorMap, N₁::Int, N₂::Int, backend, allocator)
+    TT = TensorKit.tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t))
+    S = spacetype(t)
+    Nout, Nin = numout(t), numin(t)
+    V = ProductSpace{S, Nout}(BlockTensorKit.oplus.(codomain(t).spaces)) ←
+        ProductSpace{S, Nin}(BlockTensorKit.oplus.(domain(t).spaces))
+
+    tdense = TensorOperations.tensoralloc(TT, V, Val(true), allocator)
+    BlockTensorKit.issparse(t) && zerovector!(tdense)
+    BlockTensorKit._copy_subblocks!(tdense, t)
+
+    env = repartition(fuse_legs(tdense, N₁, N₂), 2, 2; copy = true, backend, allocator)
+    TensorOperations.tensorfree!(tdense, allocator)
+
+    return env
+end
+
+function _prepare_GL_O(GL, O, backend, allocator)
+    cp = allocator_checkpoint!(allocator)
+
+    TC = TensorOperations.promote_contract(scalartype(GL), scalartype(O))
+    GL_O = TensorOperations.tensoralloc_contract(
+        TC,
+        GL, ((1, 3), (2,)), false,
+        O, ((1,), (2, 3, 4)), false,
+        ((1, 3, 5), (2, 4)), Val(true), allocator
+    )
+    @plansor backend = backend allocator = allocator begin
+        GL_O[-1 -2 -3; -4 -5] = GL[-1 1; -4] * O[1 -2; -5 -3]
+    end
+    leftenv = _fuse_env(GL_O, 1, 2, backend, allocator)
+
+    TensorOperations.tensorfree!(GL_O, allocator)
+    allocator_reset!(allocator, cp)
+
+    return leftenv
+end
+
+function _prepare_O_GR(O, GR, backend, allocator)
+    cp = allocator_checkpoint!(allocator)
+
+    TC = TensorOperations.promote_contract(scalartype(O), scalartype(GR))
+    O_GR = TensorOperations.tensoralloc_contract(
+        TC,
+        O, ((1, 2, 3), (4,)), false,
+        GR, ((2,), (1, 3)), false,
+        ((4, 3), (5, 2, 1)), Val(true), allocator
+    )
     @plansor backend = backend allocator = allocator begin
-        GL_O[-1 -2 -3; -4 -5] := H.leftenv[-1 1; -4] * H.operators[1][1 -2; -5 -3]
+        O_GR[-1 -2; -4 -5 -3] = O[-3 -5; -2 1] * GR[-1 1; -4]
     end
-    leftenv = GL_O isa TensorMap ? GL_O : TensorMap(GL_O)
-    leftenv = repartition(fuse_legs(leftenv, 1, 2), 2, 2)
+    rightenv = _fuse_env(O_GR, 2, 1, backend, allocator)
+
+    TensorOperations.tensorfree!(O_GR, allocator)
+    allocator_reset!(allocator, cp)
+
+    return rightenv
+end
+
+function prepare_operator!!(H::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor})
+    backend, allocator = H.backend, H.allocator
+
+    GL, O = H.leftenv, H.operators[1]
+    leftenv = _prepare_GL_O(GL, O, backend, allocator)
+
     rightenv = H.rightenv isa TensorMap ? H.rightenv : TensorMap(H.rightenv)
 
     return prepared_operator_type(typeof(H))(leftenv, rightenv, backend, allocator)
@@ -236,15 +308,11 @@ function prepare_operator!!(
         H::MPO_AC2_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPOTensor, <:MPSTensor}
     )
     backend, allocator = H.backend, H.allocator
-    @plansor backend = backend allocator = allocator begin
-        GL_O[-1 -2 -3; -4 -5] := H.leftenv[-1 1; -4] * H.operators[1][1 -2; -5 -3]
-        O_GR[-1 -2; -4 -5 -3] := H.operators[2][-3 -5; -2 1] * H.rightenv[-1 1; -4]
-    end
-    leftenv = GL_O isa TensorMap ? GL_O : TensorMap(GL_O)
-    leftenv = repartition(fuse_legs(leftenv, 1, 2), 2, 2)
 
-    rightenv = O_GR isa TensorMap ? O_GR : TensorMap(O_GR)
-    rightenv = repartition(fuse_legs(rightenv, 2, 1), 2, 2)
+    GL, O₁, O₂, GR = H.leftenv, H.operators[1], H.operators[2], H.rightenv
+    leftenv = _prepare_GL_O(GL, O₁, backend, allocator)
+    rightenv = _prepare_O_GR(O₂, GR, backend, allocator)
+
     return prepared_operator_type(typeof(H))(leftenv, rightenv, backend, allocator)
 end
 

src/algorithms/excitation/exci_transfer_system.jl

@@ -1,6 +1,7 @@
 function left_excitation_transfer_system(
         GBL, H::InfiniteMPOHamiltonian, exci;
-        mom = exci.momentum, solver = Defaults.linearsolver
+        mom = exci.momentum, solver = Defaults.linearsolver,
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
     )
     len = length(H)
     found = zerovector(GBL)
@@ -17,7 +18,7 @@ function left_excitation_transfer_system(
         # this would require to check the finite state machine, and discard non-connected
         # terms.
         H_partial = map(h -> getindex(h, 1:i, 1, 1, 1:i), parent(H))
-        T = TransferMatrix(exci.right_gs.AR, H_partial, exci.left_gs.AL)
+        T = TransferMatrix(exci.right_gs.AR, H_partial, exci.left_gs.AL; backend, allocator)
         start = scale!(last(found[1:i] * T), cis(-mom * len))
         if istrivial(exci) && isidentitylevel(H, i)
             regularize!(start, ρ_right, ρ_left)
@@ -27,13 +28,14 @@ function left_excitation_transfer_system(
 
         if !isemptylevel(H, i)
             if isidentitylevel(H, i)
-                T = TransferMatrix(exci.right_gs.AR, exci.left_gs.AL)
+                T = TransferMatrix(exci.right_gs.AR, exci.left_gs.AL; backend, allocator)
                 if istrivial(exci)
                     T = regularize(T, ρ_left, ρ_right)
                 end
             else
                 T = TransferMatrix(
-                    exci.right_gs.AR, map(h -> h[i, 1, 1, i], parent(H)), exci.left_gs.AL
+                    exci.right_gs.AR, map(h -> h[i, 1, 1, i], parent(H)), exci.left_gs.AL;
+                    backend, allocator
                 )
             end
 
@@ -49,8 +51,8 @@ end
 
 function right_excitation_transfer_system(
         GBR, H::InfiniteMPOHamiltonian, exci;
-        mom = exci.momentum,
-        solver = Defaults.linearsolver
+        mom = exci.momentum, solver = Defaults.linearsolver,
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
     )
     len = length(H)
     found = zerovector(GBR)
@@ -67,7 +69,7 @@ function right_excitation_transfer_system(
         # this would require to check the finite state machine, and discard non-connected
         # terms.
         H_partial = map(h -> h[i:end, 1, 1, i:end], parent(H))
-        T = TransferMatrix(exci.left_gs.AL, H_partial, exci.right_gs.AR)
+        T = TransferMatrix(exci.left_gs.AL, H_partial, exci.right_gs.AR; backend, allocator)
         start = scale!(first(T * found[i:odim]), cis(mom * len))
         if istrivial(exci) && isidentitylevel(H, i)
             regularize!(start, ρ_left, ρ_right)
@@ -77,13 +79,14 @@ function right_excitation_transfer_system(
 
         if !isemptylevel(H, i)
             if isidentitylevel(H, i)
-                tm = TransferMatrix(exci.left_gs.AL, exci.right_gs.AR)
+                tm = TransferMatrix(exci.left_gs.AL, exci.right_gs.AR; backend, allocator)
                 if istrivial(exci)
                     tm = regularize(tm, ρ_left, ρ_right)
                 end
             else
                 tm = TransferMatrix(
-                    exci.left_gs.AL, map(h -> h[i, 1, 1, i], parent(H)), exci.right_gs.AR
+                    exci.left_gs.AL, map(h -> h[i, 1, 1, i], parent(H)), exci.right_gs.AR;
+                    backend, allocator
                 )
             end
 

src/algorithms/excitation/quasiparticleexcitation.jl

@@ -28,19 +28,25 @@ Used as the `algorithm` argument of [`excitations`](@ref).
 
 * [Haegeman et al. Phys. Rev. Let. 111 (2013)](@cite haegeman2013)
 """
-struct QuasiparticleAnsatz{A, E} <: Algorithm
+struct QuasiparticleAnsatz{A, E, B} <: Algorithm
     "algorithm used for the eigenvalue solvers"
     alg::A
 
     "algorithm used for the quasiparticle environments"
     alg_environments::E
+
+    "backend for tensor contractions and index manipulations"
+    backend::B
 end
+QuasiparticleAnsatz(alg, alg_environments) =
+    QuasiparticleAnsatz(alg, alg_environments, Defaults.backend())
 function QuasiparticleAnsatz(;
         alg_environments = Defaults.alg_environments(; dynamic_tols = false),
+        backend = Defaults.backend(),
         kwargs...
     )
     alg = Defaults.alg_eigsolve(; dynamic_tols = false, kwargs...)
-    return QuasiparticleAnsatz(alg, alg_environments)
+    return QuasiparticleAnsatz(alg, alg_environments, backend)
 end
 
 ################################################################################
@@ -51,7 +57,7 @@ function excitations(H, alg::QuasiparticleAnsatz, ϕ₀::InfiniteQP, lenvs, renv
     E = effective_excitation_renormalization_energy(H, ϕ₀, lenvs, renvs)
     H_eff = EffectiveExcitationHamiltonian(H, lenvs, renvs, E)
     Es, ϕs, convhist = eigsolve(ϕ₀, num, :SR, alg.alg) do ϕ
-        return H_eff(ϕ, alg.alg_environments)
+        return H_eff(ϕ, alg.alg_environments; alg.backend)
     end
     convhist.converged < num &&
         @warn "excitation failed to converge: normres = $(convhist.normres)"
@@ -157,7 +163,7 @@ function excitations(
     E = effective_excitation_renormalization_energy(H, ϕ₀, lenvs, renvs)
     H_eff = EffectiveExcitationHamiltonian(H, lenvs, renvs, E)
     Es, ϕs, convhist = eigsolve(ϕ₀, num, :SR, alg.alg) do ϕ
-        return H_eff(ϕ, alg.alg_environments)
+        return H_eff(ϕ, alg.alg_environments; alg.backend)
     end
 
     convhist.converged < num &&
@@ -300,24 +306,37 @@ end
 # to allow Multiline checks
 Base.length(H::EffectiveExcitationHamiltonian) = length(H.operator)
 
-function (H::EffectiveExcitationHamiltonian)(ϕ::QP, alg_environments = DefaultAlgorithm())
-    qp_envs = environments(ϕ, H.operator, ϕ, alg_environments; lenvs = H.lenvs, renvs = H.renvs)
-    return effective_excitation_hamiltonian(H.operator, ϕ, qp_envs, H.energy)
+function (H::EffectiveExcitationHamiltonian)(
+        ϕ::QP, alg_environments = DefaultAlgorithm();
+        backend::AbstractBackend = DefaultBackend()
+    )
+    qp_envs = environments(
+        ϕ, H.operator, ϕ, alg_environments; lenvs = H.lenvs, renvs = H.renvs, backend
+    )
+    return effective_excitation_hamiltonian(H.operator, ϕ, qp_envs, H.energy; backend)
 end
 function (H::Multiline{<:EffectiveExcitationHamiltonian})(
-        ϕ::MultilineQP, alg_environments = DefaultAlgorithm()
+        ϕ::MultilineQP, alg_environments = DefaultAlgorithm(); kwargs...
     )
-    return Multiline(map((x, y) -> x(y, alg_environments), parent(H), parent(ϕ)))
+    return Multiline(map((x, y) -> x(y, alg_environments; kwargs...), parent(H), parent(ϕ)))
 end
 
 function effective_excitation_hamiltonian(H, ϕ, envs = environments(ϕ, H))
     E₀ = effective_excitation_renormalization_energy(H, ϕ, envs.leftenvs, envs.rightenvs)
     return effective_excitation_hamiltonian(H, ϕ, envs, E₀)
 end
-function effective_excitation_hamiltonian(H, ϕ, qp_envs, E)
+function effective_excitation_hamiltonian(
+        H, ϕ, qp_envs, E, scheduler = Defaults.scheduler[];
+        backend::AbstractBackend = DefaultBackend()
+    )
     ϕ′ = similar(ϕ)
-    tforeach(1:length(ϕ); scheduler = Defaults.scheduler[]) do loc
-        ϕ′[loc] = _effective_excitation_local_apply(loc, ϕ, H, E[loc], qp_envs)
+    # This is the site that fans out, so it is the one that reads the scheduler and derives
+    # the allocator from it
+    allocator = default_allocator(ϕ.left_gs, scheduler)
+    tforeach(1:length(ϕ); scheduler) do loc
+        ϕ′[loc] = _effective_excitation_local_apply(
+            loc, ϕ, H, E[loc], qp_envs; backend, allocator
+        )
         return nothing
     end
     return ϕ′
@@ -338,10 +357,13 @@ function effective_excitation_hamiltonian(H::MultilineMPO, ϕ::MultilineQP, envs
     )
 end
 
-function _effective_excitation_local_apply(site, ϕ, H::MPOHamiltonian, E::Number, envs)
+function _effective_excitation_local_apply(
+        site, ϕ, H::MPOHamiltonian, E::Number, envs;
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
+    )
     B = ϕ[site]
-    GL = leftenv(envs.leftenvs, site, ϕ.left_gs)
-    GR = rightenv(envs.rightenvs, site, ϕ.right_gs)
+    GL = leftenv(envs.leftenvs, site, ϕ.left_gs; backend, allocator)
+    GR = rightenv(envs.rightenvs, site, ϕ.right_gs; backend, allocator)
 
     # renormalize first -> allocates destination
     B′ = scale(B, -E)
@@ -366,13 +388,16 @@ function _effective_excitation_local_apply(site, ϕ, H::MPOHamiltonian, E::Numbe
     return B′
 end
 
-function _effective_excitation_local_apply(site, ϕ, H::MPO, E::Number, envs)
+function _effective_excitation_local_apply(
+        site, ϕ, H::MPO, E::Number, envs;
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
+    )
     left_gs = ϕ.left_gs
     right_gs = ϕ.right_gs
 
     B = ϕ[site]
-    GL = leftenv(envs.leftenvs, site, ϕ.left_gs)
-    GR = rightenv(envs.rightenvs, site, ϕ.right_gs)
+    GL = leftenv(envs.leftenvs, site, ϕ.left_gs; backend, allocator)
+    GR = rightenv(envs.rightenvs, site, ϕ.right_gs; backend, allocator)
 
     @plansor T[-1 -2; -3 -4] := GL[-1 5; 4] * B[4 2; -3 1] * H[site][5 -2; 2 3] * GR[1 3; -4]