diff --git a/Project.toml b/Project.toml index 0882fa3a4..1e8717e06 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,7 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36" MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" @@ -46,12 +47,13 @@ LoggingExtras = "~1.0" MatrixAlgebraKit = "0.6.5" OhMyThreads = "0.7, 0.8" OptimKit = "0.3.1, 0.4, 0.5" +Preferences = "1" Printf = "1" Random = "1" RecipesBase = "1.1" TensorKit = "0.17.1" TensorKitManifolds = "0.7, 0.8" -TensorOperations = "5.5.1" +TensorOperations = "5.6.2" TimerOutputs = "0.5.29" VectorInterface = "0.2, 0.3, 0.4, 0.5, 0.6" julia = "1.10" diff --git a/benchmark/gpu-allocator/run_gpu_tests.sbatch b/benchmark/gpu-allocator/run_gpu_tests.sbatch new file mode 100644 index 000000000..4a2966565 --- /dev/null +++ b/benchmark/gpu-allocator/run_gpu_tests.sbatch @@ -0,0 +1,63 @@ +#!/bin/bash +# Run the MPSKit CUDA test suite, including the new algorithm-level allocator tests. +# +# There is no GPU on the workstation, so this is the only way to exercise +# `MPSKit.allocatortype(::Type{<:CuArray})` and confirm that a device-backed DMRG/TDVP draws its +# scratch space from `CUDAAllocator` rather than a host buffer. +# +# Submit from the repository root: +# sbatch benchmark/gpu-allocator/run_gpu_tests.sbatch +# +# Adjust --partition/--gres to whatever GPU partition you have access to; the values below are +# placeholders. Check with `sinfo -o "%P %G %N"`. + +#SBATCH --job-name=mpskit-gpu-alloc +#SBATCH --output=benchmark/gpu-allocator/logs/%x-%j.out +#SBATCH --error=benchmark/gpu-allocator/logs/%x-%j.err +#SBATCH --time=01:00:00 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --partition=gpu # placeholder +#SBATCH --gres=gpu:1 # placeholder + +set -euo pipefail + +cd "${SLURM_SUBMIT_DIR:?must be submitted with sbatch}" +mkdir -p benchmark/gpu-allocator/logs + +# The depot is precompiled under 1.12.6; `module load julia` gives 1.11.2, which fails with a +# StaticData error. Use the juliaup install on GPFS directly. +JULIA="${HOME}/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/bin/julia" +test -x "$JULIA" || { echo "julia not found at $JULIA" >&2; exit 1; } + +echo "host: $(hostname)" +echo "julia: $("$JULIA" --version)" +nvidia-smi || echo "WARNING: nvidia-smi unavailable" + +export JULIA_NUM_THREADS=4 +# one BLAS thread per Julia thread would oversubscribe; the GPU does the heavy lifting anyway +export OPENBLAS_NUM_THREADS=1 + +"$JULIA" --project=test -e ' + using Pkg + Pkg.instantiate() + + using CUDA, cuTENSOR + @info "CUDA status" functional = CUDA.functional() cutensor = cuTENSOR.functional() + CUDA.functional() || error("CUDA is not functional on this node") + CUDA.versioninfo() + + include(joinpath(@__DIR__, "test", "setup", "testsetup.jl")) + using .TestSetup + using Test + + @testset "MPSKit CUDA" verbose = true begin + for f in ("cu_adapt.jl", "states.jl", "operators.jl", "algorithms.jl") + @testset "$f" begin + include(joinpath(@__DIR__, "test", "gpu", "cuda", f)) + end + end + end +' diff --git a/docs/src/changelog.md b/docs/src/changelog.md index 7472c9cfb..dddbd57c6 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -32,6 +32,14 @@ When releasing a new version, move the "Unreleased" changes to a new version sec Unlike `TDVP` it has no backward-in-time substep (stable for imaginary-time evolution), and passing a truncating `trunc` enables rank-adaptivity (the bond dimension grows and shrinks automatically to track entanglement). +- A `backend` setting on every algorithm, for its tensor contractions and index manipulations, + defaulting to `MPSKit.Defaults.backend()`. + ([#467](https://github.com/QuantumKitHub/MPSKit.jl/pull/467)) +- Local updates now serve their intermediate tensors from a dedicated allocator, selected internally + by `MPSKit.default_allocator`, instead of leaving them to the garbage collector + (two-site DMRG: -64% allocations, -57% GC time, -23% wall time). + Disable with `MPSKit.Defaults.set_buffering!(false)`. + ([#467](https://github.com/QuantumKitHub/MPSKit.jl/pull/467)) ### Changed diff --git a/docs/src/man/parallelism.md b/docs/src/man/parallelism.md index 683be061c..4e15746a3 100644 --- a/docs/src/man/parallelism.md +++ b/docs/src/man/parallelism.md @@ -98,6 +98,29 @@ MPSKit.Defaults.set_scheduler!(:dynamic) # default: multithreading with some loa For further reference on the available schedulers and finer control, please refer to the [`OhMyThreads.jl` documentation](https://juliafolds2.github.io/OhMyThreads.jl/stable/) +## Allocators + +Every local update in an MPS algorithm contracts a handful of intermediate tensors that are thrown away immediately afterwards. +Rather than asking Julia's garbage collector to clean up after each one, MPSKit hands the contractions an allocator that manages those intermediates itself. +On a typical ground-state search this cuts both the allocation count and the garbage-collection time substantially. + +Which allocator is appropriate depends on two things, so [`MPSKit.default_allocator`](@ref) is asked for one at the start of every solve, and the answer is then used for all of its local updates: + +- *Where the tensors live.* Host memory can be served by an allocator that bypasses Julia's memory manager, whereas a device-backed state needs its scratch space on the device. This is only known once there is a state in hand, which is why it is not a setting on the algorithm. +- *Whether the allocator will be shared between tasks.* `TensorOperations.BufferAllocator` is a bump buffer with a mutable offset, so it is only safe when a single task owns it - a serial sweep. Concurrent work instead gets a `TensorOperations.ManualAllocator`, which `malloc`s and `free`s each intermediate and holds no state at all. + +Anything MPSKit does not recognise falls back on `TensorOperations.DefaultAllocator`, which allocates through the storage type itself: correct on any device, at the cost of leaving the intermediates to the garbage collector. +That is also what device-backed states get. +Note that the `backend` field needs no such treatment: `TensorOperations.DefaultBackend` is a placeholder that resolves to an implementation based on the types of the tensors involved, so it already selects cuTENSOR for `CuArray`-backed states. + +Dedicated scratch space can be turned off when memory rather than time is the binding constraint, which trades it back for garbage-collector pressure: + +```julia +MPSKit.Defaults.set_buffering!(false) +``` + +This is a compile-time preference: it is written to a `LocalPreferences.toml` file next to the active `Project.toml`, and Julia has to be restarted for the change to take effect. + ## TensorKit multithreading Finally, when dealing with tensors that have some internal symmetry, it is also possible to diff --git a/src/MPSKit.jl b/src/MPSKit.jl index 920a8fefa..29077215b 100644 --- a/src/MPSKit.jl +++ b/src/MPSKit.jl @@ -57,6 +57,7 @@ using Compat: @compat @compat public DynamicTols, DynamicTol, AdaptiveKrylov @compat public VERBOSE_NONE, VERBOSE_WARN, VERBOSE_CONV, VERBOSE_ITER, VERBOSE_ALL @compat public IterLog, loginit!, logiter!, logfinish!, logcancel! +@compat public default_allocator # Imports # ------- @@ -69,7 +70,7 @@ using BlockTensorKit using BlockTensorKit: TensorMapSumSpace using TensorOperations using TensorOperations: AbstractBackend, DefaultBackend, DefaultAllocator, BufferAllocator, - allocator_checkpoint!, allocator_reset! + ManualAllocator, allocator_checkpoint!, allocator_reset! using KrylovKit using KrylovKit: KrylovAlgorithm using OptimKit @@ -99,6 +100,7 @@ using .DynamicTols include("utility/defaults.jl") using .Defaults: VERBOSE_NONE, VERBOSE_WARN, VERBOSE_CONV, VERBOSE_ITER, VERBOSE_ALL +include("utility/allocator.jl") include("utility/logging.jl") using .IterativeLoggers include("utility/iterativesolvers.jl") diff --git a/src/algorithms/ED.jl b/src/algorithms/ED.jl index 7b42d6888..915bfd782 100644 --- a/src/algorithms/ED.jl +++ b/src/algorithms/ED.jl @@ -2,7 +2,8 @@ exact_diagonalization( H::FiniteMPOHamiltonian; sector = rightunit(H), num::Int = 1, which::Symbol = :SR, - alg = Defaults.alg_eigsolve(; dynamic_tols = false) + alg = Defaults.alg_eigsolve(; dynamic_tols = false), + backend = Defaults.backend() ) -> vals, state_vecs, convhist Use [`KrylovKit.eigsolve`](@extref) to perform exact diagonalization on a @@ -21,6 +22,7 @@ equivalent to dense eigenvectors. - `which::Symbol = :SR`: the kind eigenvalues to find, see [`KrylovKit.eigsolve`](@extref). - `alg = Defaults.alg_eigsolve(; dynamic_tols = false)`: the diagonalization algorithm to use, see [`KrylovKit.eigsolve`](@extref). +- `backend = Defaults.backend()`: backend for tensor contractions and index manipulations. !!! note "Valid `sector` values" The total charge of the eigenvectors is imposed by adding a charged auxiliary space as @@ -34,7 +36,8 @@ equivalent to dense eigenvectors. function exact_diagonalization( H::FiniteMPOHamiltonian; sector = rightunit(H), num::Int = 1, which::Symbol = :SR, - alg = Defaults.alg_eigsolve(; dynamic_tols = false) + alg = Defaults.alg_eigsolve(; dynamic_tols = false), + backend::AbstractBackend = Defaults.backend() ) L = length(H) @assert L > 1 "FiniteMPOHamiltonian must have length > 1" @@ -74,7 +77,8 @@ function exact_diagonalization( # optimize the middle site # Because the MPS is full rank - this is equivalent to the full Hamiltonian AC₀ = state.AC[middle_site] - H_ac = AC_hamiltonian(middle_site, state, H, state, envs) + allocator = default_allocator(state, SerialScheduler()) + H_ac = AC_hamiltonian(middle_site, state, H, state, envs; backend, allocator) vals, vecs, convhist = eigsolve(H_ac, AC₀, num, which, alg) # repack eigenstates diff --git a/src/algorithms/approximate/fvomps.jl b/src/algorithms/approximate/fvomps.jl index aacdacf2b..47dcfdfef 100644 --- a/src/algorithms/approximate/fvomps.jl +++ b/src/algorithms/approximate/fvomps.jl @@ -1,4 +1,5 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environments(ψ, _environment_args(Oϕ)...)) + allocator = default_allocator(ψ, SerialScheduler()) ϵ::Float64 = 2 * alg.tol log = IterLog("DMRG2") @@ -7,7 +8,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environment for iter in 1:(alg.maxiter) ϵ = 0.0 for pos in [1:(length(ψ) - 1); (length(ψ) - 2):-1:1] - AC2′ = AC2_projection(pos, ψ, Oϕ, envs) + AC2′ = AC2_projection(pos, ψ, Oϕ, envs; alg.backend, allocator) al, c, ar, = svd_trunc!(AC2′, inner_alg_gauge(alg)) AC2 = ψ.AC[pos] * _transpose_tail(ψ.AR[pos + 1]) @@ -36,6 +37,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environment end function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG, envs = environments(ψ, _environment_args(Oϕ)...)) + allocator = default_allocator(ψ, SerialScheduler()) ϵ::Float64 = 2 * alg.tol log = IterLog("DMRG") @@ -44,7 +46,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG, envs = environments for iter in 1:(alg.maxiter) ϵ = 0.0 for pos in [1:(length(ψ) - 1); length(ψ):-1:2] - AC′ = AC_projection(pos, ψ, Oϕ, envs) + AC′ = AC_projection(pos, ψ, Oϕ, envs; alg.backend, allocator) AC = ψ.AC[pos] ϵ = max(ϵ, norm(AC′ - AC) / norm(AC′)) diff --git a/src/algorithms/approximate/idmrg.jl b/src/algorithms/approximate/idmrg.jl index ea2475860..7503ac76d 100644 --- a/src/algorithms/approximate/idmrg.jl +++ b/src/algorithms/approximate/idmrg.jl @@ -2,6 +2,7 @@ function approximate!( ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS}, alg::IDMRG, envs = environments(ψ, toapprox...) ) + allocator = default_allocator(ψ, SerialScheduler()) log = IterLog("IDMRG") ϵ::Float64 = 2 * alg.tol local iter @@ -15,7 +16,8 @@ function approximate!( for col in 1:size(ψ, 2) for row in 1:size(ψ, 1) ψ.AC[row + 1, col] = AC_projection( - CartesianIndex(row, col), ψ, toapprox, envs + CartesianIndex(row, col), ψ, toapprox, envs; + alg.backend, allocator ) normalize!(ψ.AC[row + 1, col]) ψ.AL[row + 1, col], ψ.C[row + 1, col] = left_orth!(ψ.AC[row + 1, col]) @@ -27,7 +29,8 @@ function approximate!( for col in reverse(1:size(ψ, 2)) for row in 1:size(ψ, 1) ψ.AC[row + 1, col] = AC_projection( - CartesianIndex(row, col), ψ, toapprox, envs + CartesianIndex(row, col), ψ, toapprox, envs; + alg.backend, allocator ) normalize!(ψ.AC[row + 1, col]) ψ.C[row + 1, col - 1], temp = right_orth!(_transpose_tail(ψ.AC[row + 1, col])) @@ -64,6 +67,7 @@ function approximate!( ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS}, alg::IDMRG2, envs = environments(ψ, toapprox...) ) + allocator = default_allocator(ψ, SerialScheduler()) size(ψ, 2) < 2 && throw(ArgumentError("unit cell should be >= 2")) ϵ::Float64 = 2 * alg.tol log = IterLog("IDMRG2") @@ -80,7 +84,7 @@ function approximate!( for row in 1:size(ψ, 1) AC2′ = AC2_projection( CartesianIndex(row, site), ψ, toapprox, envs; - kind = :ACAR + kind = :ACAR, alg.backend, allocator ) al, c, ar = svd_trunc!(AC2′; trunc = alg.trunc, alg = alg.alg_svd) normalize!(c) @@ -99,7 +103,10 @@ function approximate!( ψ.AL[1, end] = ψ.AC[1, end] / ψ.C[1, end] ψ.AC[1, 1] = _mul_tail(ψ.AL[1, 1], ψ.C[1, 1]) for row in 1:size(ψ, 1) - AC2′ = AC2_projection(CartesianIndex(row, size(ψ, 2)), ψ, toapprox, envs; kind = :ALAC) + AC2′ = AC2_projection( + CartesianIndex(row, size(ψ, 2)), ψ, toapprox, envs; + kind = :ALAC, alg.backend, allocator + ) al, c, ar = svd_trunc!(AC2′; trunc = alg.trunc, alg = alg.alg_svd) normalize!(c) @@ -123,7 +130,7 @@ function approximate!( for row in 1:size(ψ, 1) AC2′ = AC2_projection( CartesianIndex(row, site), ψ, toapprox, envs; - kind = :ALAC + kind = :ALAC, alg.backend, allocator ) al, c, ar = svd_trunc!(AC2′; trunc = alg.trunc, alg = alg.alg_svd) normalize!(c) @@ -141,7 +148,10 @@ function approximate!( ψ.AC[1, end] = _mul_front(ψ.C[1, end - 1], ψ.AR[1, end]) ψ.AR[1, 1] = _transpose_front(ψ.C[1, end] \ _transpose_tail(ψ.AC[1, 1])) for row in 1:size(ψ, 1) - AC2′ = AC2_projection(CartesianIndex(row, 0), ψ, toapprox, envs; kind = :ACAR) + AC2′ = AC2_projection( + CartesianIndex(row, 0), ψ, toapprox, envs; + kind = :ACAR, alg.backend, allocator + ) al, c, ar = svd_trunc!(AC2′; trunc = alg.trunc, alg = alg.alg_svd) normalize!(c) @@ -155,7 +165,6 @@ function approximate!( transfer_leftenv!(envs, ψ, toapprox, 1) transfer_rightenv!(envs, ψ, toapprox, 0) - normalize!(envs, ψ, toapprox) # update error diff --git a/src/algorithms/approximate/vomps.jl b/src/algorithms/approximate/vomps.jl index 70f3335b0..77e6c8cde 100644 --- a/src/algorithms/approximate/vomps.jl +++ b/src/algorithms/approximate/vomps.jl @@ -5,7 +5,10 @@ Base.@deprecate( ), approximate( ψ, toapprox, - VOMPS(; alg.tol, alg.maxiter, alg.finalize, alg.verbosity, alg.alg_gauge, alg.alg_environments), + VOMPS(; + alg.tol, alg.maxiter, alg.finalize, alg.verbosity, alg.alg_gauge, + alg.alg_environments, alg.backend, + ), envs...; kwargs... ) ) @@ -14,9 +17,13 @@ function approximate( mps::MultilineMPS, toapprox::Tuple{<:MultilineMPO, <:MultilineMPS}, alg::VOMPS, envs = environments(mps, toapprox...) ) + return _approximate_vomps(mps, toapprox, alg, envs) +end + +function _approximate_vomps(mps, toapprox, alg::VOMPS, envs) log = IterLog("VOMPS") iter = 0 - ϵ = calc_galerkin(mps, toapprox..., envs) + ϵ = calc_galerkin(mps, toapprox..., envs; alg.backend) alg_environments = adapt_solver(alg.alg_environments; iter, g_global = ϵ) recalculate!(envs, mps, toapprox..., alg_environments) @@ -54,7 +61,7 @@ function Base.iterate(it::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any, <:T mps, envs = it.finalize(state.iter, mps, state.operator, envs)::typeof((mps, envs)) # error criterion - ϵ = calc_galerkin(mps, state.operator..., envs) + ϵ = calc_galerkin(mps, state.operator..., envs; it.backend) # update state it.state = VOMPSState(mps, state.operator, envs, state.iter + 1, ϵ) @@ -71,17 +78,24 @@ function localupdate_step!( ACs = similar(state.mps.AC) dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs - foreach(eachsite(state.mps)) do site + # the sweep is serial, so a single allocator serves all sites + allocator = default_allocator(state.mps, SerialScheduler()) + for site in eachsite(state.mps) AC = map(1:size(state.mps, 1)) do row - AC_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs) + AC_projection( + CartesianIndex(row, site), state.mps, state.operator, state.envs; + it.backend, allocator + ) end circshift!(AC, 1) C = map(1:size(state.mps, 1)) do row - C_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs) + C_projection( + CartesianIndex(row, site), state.mps, state.operator, state.envs; + it.backend, allocator + ) end circshift!(C, 1) dst_ACs[site] = regauge!(AC, C; alg = alg_orth) - return nothing end return ACs @@ -95,18 +109,27 @@ function localupdate_step!( ACs = similar(state.mps.AC) dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs + # every site - and the AC and C projections within a site - runs concurrently, so the allocator + # is shared and has to be one that tolerates that + allocator = default_allocator(state.mps, scheduler) tforeach(eachsite(state.mps); scheduler) do site local AC, C @sync begin Threads.@spawn begin AC = map(1:size(state.mps, 1)) do row - AC_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs) + AC_projection( + CartesianIndex(row, site), state.mps, state.operator, state.envs; + it.backend, allocator + ) end circshift!(AC, 1) end Threads.@spawn begin C = map(1:size(state.mps, 1)) do row - C_projection(CartesianIndex(row, site), state.mps, state.operator, state.envs) + C_projection( + CartesianIndex(row, site), state.mps, state.operator, state.envs; + it.backend, allocator + ) end circshift!(C, 1) end diff --git a/src/algorithms/changebonds/optimalexpand.jl b/src/algorithms/changebonds/optimalexpand.jl index ee23d69ea..0a039ddf3 100644 --- a/src/algorithms/changebonds/optimalexpand.jl +++ b/src/algorithms/changebonds/optimalexpand.jl @@ -34,12 +34,15 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`changebonds`](@ref) and [`changebonds!`](@ref). """ -@kwdef struct OptimalExpand{S} <: Algorithm +@kwdef struct OptimalExpand{S, B} <: Algorithm "algorithm used for the singular value decomposition" alg_svd::S = Defaults.alg_svd() "[truncation strategy](@extref MatrixAlgebraKit.TruncationStrategy) selecting how many directions are *added* to each bond, rather than how much of the bond is kept" trunc::TruncationStrategy + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end # Simple wrapper to convert between diffrent type of InifniteMPS. @@ -56,12 +59,13 @@ function changebonds( ψ::InfiniteMPS, H::InfiniteMPOHamiltonian, alg::OptimalExpand, envs = environments(ψ, H, ψ) ) + allocator = default_allocator(ψ, SerialScheduler()) T = eltype(ψ.AL) AL′ = similar(ψ.AL) AR′ = similar(ψ.AR, tensormaptype(spacetype(T), 1, numind(T) - 1, storagetype(T))) for i in 1:length(ψ) # determine optimal expansion spaces around bond i - AC2 = AC2_projection(i, ψ, H, ψ, envs; kind = :ACAR) + AC2 = AC2_projection(i, ψ, H, ψ, envs; kind = :ACAR, alg.backend, allocator) # Use the nullspaces and SVD decomposition to determine the optimal expansion space VL = left_null(ψ.AL[i]) @@ -82,6 +86,7 @@ function changebonds( end function changebonds(ψ::MultilineMPS, H, alg::OptimalExpand, envs = environments(ψ, H, ψ)) + allocator = default_allocator(ψ, SerialScheduler()) TL = eltype(ψ.AL) AL′ = PeriodicMatrix{TL}(undef, size(ψ.AL)) TR = tensormaptype(spacetype(TL), 1, numind(TL) - 1, storagetype(TL)) @@ -89,7 +94,10 @@ function changebonds(ψ::MultilineMPS, H, alg::OptimalExpand, envs = environment # determine optimal expansion spaces around bond i for i in 1:size(ψ, 1), j in 1:size(ψ, 2) - AC2 = AC2_projection(CartesianIndex(i - 1, j), ψ, H, ψ, envs; kind = :ACAR) + AC2 = AC2_projection( + CartesianIndex(i - 1, j), ψ, H, ψ, envs; + kind = :ACAR, alg.backend, allocator + ) # Use the nullspaces and SVD decomposition to determine the optimal expansion space VL = left_null(ψ.AL[i, j]) @@ -108,10 +116,14 @@ function changebonds(ψ::MultilineMPS, H, alg::OptimalExpand, envs = environment return newψ, envs end - # Finite system # ------------- -function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::OptimalExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::OptimalExpand, envs; + normalize::Bool = true, + # a sweep that already holds one should pass it: this is called once per site + allocator = default_allocator(ψ, SerialScheduler()), + ) bond = site left = ψ.AC[site] right = ψ.AR[site + 1] @@ -119,7 +131,7 @@ function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::Op NR = right_null!(_transpose_tail(right; copy = true)) # two-site update from the projected effective Hamiltonian - AC2 = AC2_projection(bond, ψ, H, ψ, envs) + AC2 = AC2_projection(bond, ψ, H, ψ, envs; alg.backend, allocator) # select the dominant directions in the complement of the current state g2 = adjoint(NL) * AC2 * adjoint(NR) @@ -141,7 +153,12 @@ function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::Op ψ.AC[site + 1] = (nc, nar) return ψ end -function changebond!(site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::OptimalExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::OptimalExpand, envs; + normalize::Bool = true, + # a sweep that already holds one should pass it: this is called once per site + allocator = default_allocator(ψ, SerialScheduler()), + ) bond = site - 1 left = ψ.AL[site - 1] right = ψ.AC[site] @@ -149,7 +166,7 @@ function changebond!(site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::Opt NR = right_null!(_transpose_tail(right; copy = true)) # two-site update from the projected effective Hamiltonian - AC2 = AC2_projection(bond, ψ, H, ψ, envs) + AC2 = AC2_projection(bond, ψ, H, ψ, envs; alg.backend, allocator) # select the dominant directions in the complement of the current state g2 = adjoint(NL) * AC2 * adjoint(NR) diff --git a/src/algorithms/changebonds/randexpand.jl b/src/algorithms/changebonds/randexpand.jl index ac948042f..fad974892 100644 --- a/src/algorithms/changebonds/randexpand.jl +++ b/src/algorithms/changebonds/randexpand.jl @@ -86,7 +86,13 @@ function changebonds!(ψ::AbstractFiniteMPS, alg::RandExpand) return normalize!(ψ) end -function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::RandExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::RandExpand, envs; + normalize::Bool = true, + # a randomized expansion has no effective-operator contraction to serve: accepted and + # dropped, so that a sweep can pass its allocator to any `alg_expand` + allocator = nothing, + ) bond = site left = ψ.AC[site] right = ψ.AR[site + 1] @@ -115,7 +121,13 @@ function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::Ra ψ.AC[site + 1] = (nc, nar) return ψ end -function changebond!(site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::RandExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::RandExpand, envs; + normalize::Bool = true, + # a randomized expansion has no effective-operator contraction to serve: accepted and + # dropped, so that a sweep can pass its allocator to any `alg_expand` + allocator = nothing, + ) bond = site - 1 left = ψ.AL[site - 1] right = ψ.AC[site] diff --git a/src/algorithms/changebonds/sketchedexpand.jl b/src/algorithms/changebonds/sketchedexpand.jl index b4d4680c9..ff57653de 100644 --- a/src/algorithms/changebonds/sketchedexpand.jl +++ b/src/algorithms/changebonds/sketchedexpand.jl @@ -38,7 +38,7 @@ Used as the `algorithm` argument of [`changebonds`](@ref) and [`changebonds!`](@ * [Gleis et al. Phys. Rev. Lett. 130, 246402 (2023)](@cite gleis2023) """ -@kwdef struct SketchedExpand{S} <: Algorithm +@kwdef struct SketchedExpand{S, B} <: Algorithm "algorithm used to orthonormalize the sketched complement (passed as the `alg` of `left_orth!`/`right_orth!`); `nothing` selects QR without oversampling and an SVD-based decomposition otherwise" alg_orth::S = nothing @@ -47,6 +47,9 @@ Used as the `algorithm` argument of [`changebonds`](@ref) and [`changebonds!`](@ "number of extra sketch columns drawn beyond the target rank (range-finder oversampling)" oversampling::Int = 0 + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end """ @@ -69,7 +72,12 @@ end # into the effective environment, and the dominant directions are read off a small SVD of the # sketched gradient. The complement projectors act with the isometric MPS tensors directly, which # leaves the MPS gauge (and any incrementally-maintained environments) untouched. -function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::SketchedExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::SketchedExpand, envs; + normalize::Bool = true, + # a sweep that already holds one should pass it: this is called once per site + allocator = default_allocator(ψ, SerialScheduler()), + ) left = ψ.AC[site] right = ψ.AR[site + 1] AL, _, _ = left_gauge(left) # local left-isometric form @@ -85,7 +93,9 @@ function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::Sk Ω = randisometry(scalartype(left), codomain(AL) ← Vℓ) Q, _ = qr_compact!(project_complement!(Ω, AL)) GL = leftenv(envs, site, ψ) * TransferMatrix(left, H[site], Q) - Hac = MPO_AC_Hamiltonian(GL, H[site + 1], rightenv(envs, site + 1, ψ)) + Hac = MPO_AC_Hamiltonian( + GL, H[site + 1], rightenv(envs, site + 1, ψ), alg.backend, allocator + ) Y = Hac * right # orthonormalize the sketched right complement, truncating away the oversampling padding @@ -108,7 +118,12 @@ function changebond!(site::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::Sk ψ.AC[site + 1] = (nc, nar) return ψ end -function changebond!(site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::SketchedExpand, envs; normalize::Bool = true) +function changebond!( + site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::SketchedExpand, envs; + normalize::Bool = true, + # a sweep that already holds one should pass it: this is called once per site + allocator = default_allocator(ψ, SerialScheduler()), + ) left = ψ.AL[site - 1] right = ψ.AC[site] _, ARtt = right_orth!(_transpose_tail(right; copy = true); trunc = notrunc()) # local right-isometric form @@ -124,7 +139,9 @@ function changebond!(site::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::Ske _, Qr_o = lq_compact!(project_complement_right!(Ω, ARtt)) Qr = _transpose_front(Qr_o) GR = TransferMatrix(right, H[site], Qr) * rightenv(envs, site, ψ) - Hac = MPO_AC_Hamiltonian(leftenv(envs, site - 1, ψ), H[site - 1], GR) + Hac = MPO_AC_Hamiltonian( + leftenv(envs, site - 1, ψ), H[site - 1], GR, alg.backend, allocator + ) Y = Hac * left # orthonormalize the sketched left complement, truncating away the oversampling padding diff --git a/src/algorithms/changebonds/vumpssvd.jl b/src/algorithms/changebonds/vumpssvd.jl index 4064c68b7..ec4c32aa8 100644 --- a/src/algorithms/changebonds/vumpssvd.jl +++ b/src/algorithms/changebonds/vumpssvd.jl @@ -14,7 +14,7 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`changebonds`](@ref). """ -@kwdef struct VUMPSSvdCut <: Algorithm +@kwdef struct VUMPSSvdCut{B} <: Algorithm "algorithm used for gauging the `InfiniteMPS`" alg_gauge = Defaults.alg_gauge(; dynamic_tols = false) @@ -26,6 +26,9 @@ Used as the `algorithm` argument of [`changebonds`](@ref). "algorithm used for [truncation](@extref MatrixAlgebraKit.TruncationStrategy) of the two-site update" trunc::TruncationStrategy + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end function changebonds_1( @@ -56,13 +59,14 @@ function changebonds_1( end function changebonds_n(state::InfiniteMPS, H, alg::VUMPSSvdCut, envs = environments(state, H, state)) + allocator = default_allocator(state, SerialScheduler()) for loc in 1:length(state) @plansor AC2[-1 -2; -3 -4] := state.AC[loc][-1 -2; 1] * state.AR[loc + 1][1 -4; -3] - Hac2 = AC2_hamiltonian(loc, state, H, state, envs) + Hac2 = AC2_hamiltonian(loc, state, H, state, envs; alg.backend, allocator) _, nAC2 = fixedpoint(Hac2, AC2, :SR, alg.alg_eigsolve) - Hc = C_hamiltonian(loc + 1, state, H, state, envs) + Hc = C_hamiltonian(loc + 1, state, H, state, envs; alg.backend, allocator) _, nC2 = fixedpoint(Hc, state.C[loc + 1], :SR, alg.alg_eigsolve) #svd ac2, get new AL1 and S,V ---> AC diff --git a/src/algorithms/derivatives/derivatives.jl b/src/algorithms/derivatives/derivatives.jl index 7d5cc677e..6c9f58c02 100644 --- a/src/algorithms/derivatives/derivatives.jl +++ b/src/algorithms/derivatives/derivatives.jl @@ -191,17 +191,28 @@ for kind in (:C, :AC, :AC2) end end -function C_projection(site, below, operator, above, envs) +function C_projection( + site, below, operator, above, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) C = above isa Multiline ? above.C[:, site] : above.C[site] - return C_hamiltonian(site, below, operator, above, envs; prepare = false) * C + H = C_hamiltonian(site, below, operator, above, envs; prepare = false, backend, allocator) + return H * C end -function AC_projection(site, below, operator, above, envs) +function AC_projection( + site, below, operator, above, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) AC = above isa Multiline ? above.AC[:, site] : above.AC[site] - return AC_hamiltonian(site, below, operator, above, envs; prepare = false) * AC + H = AC_hamiltonian(site, below, operator, above, envs; prepare = false, backend, allocator) + return H * AC end -function AC2_projection(site::Int, below, operator, above, envs; kwargs...) - return AC2_hamiltonian(site, below, operator, above, envs; prepare = false) * - AC2(above, site; kwargs...) +function AC2_projection( + site::Int, below, operator, above, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator(), kwargs... + ) + H = AC2_hamiltonian(site, below, operator, above, envs; prepare = false, backend, allocator) + return H * AC2(above, site; kwargs...) end # Multiline @@ -229,4 +240,4 @@ Given an operator, try to construct a more efficient representation of that oper This typically consists of precomputing some parts of the application, and is expected to only pay off for repeated applications. """ -prepare_operator!!(O, backend::AbstractBackend = DefaultBackend(), allocator = BufferAllocator()) = O +prepare_operator!!(O, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) = O diff --git a/src/algorithms/derivatives/hamiltonian_derivatives.jl b/src/algorithms/derivatives/hamiltonian_derivatives.jl index f965db9ac..cfe05ac21 100644 --- a/src/algorithms/derivatives/hamiltonian_derivatives.jl +++ b/src/algorithms/derivatives/hamiltonian_derivatives.jl @@ -12,42 +12,52 @@ const _HAM_MPS_TYPES = Union{ Efficient operator for representing the single-site derivative of a `MPOHamiltonian` sandwiched between two MPSs. In particular, this operator aims to make maximal use of the structure of the `MPOHamiltonian` to reduce the number of operations required to apply the operator to a tensor. """ -struct JordanMPO_AC_Hamiltonian{O1, O2, O3} <: DerivativeOperator +struct JordanMPO_AC_Hamiltonian{O1, O2, O3, Bk <: AbstractBackend, Al} <: DerivativeOperator D::Union{O1, Missing} # onsite I::Union{O1, Missing} # not started E::Union{O1, Missing} # finished C::Union{O2, Missing} # starting B::Union{O2, Missing} # ending A::Union{O3, Missing} # continuing + backend::Bk # contraction backend used by the matvec + allocator::Al # scratch-buffer allocator used by the matvec - function JordanMPO_AC_Hamiltonian{O1, O2, O3}( + function JordanMPO_AC_Hamiltonian{O1, O2, O3, Bk, Al}( D::Union{O1, Missing}, I::Union{O1, Missing}, E::Union{O1, Missing}, - C::Union{O2, Missing}, B::Union{O2, Missing}, A::Union{O3, Missing} - ) where {O1, O2, O3} - return new{O1, O2, O3}(D, I, E, C, B, A) + C::Union{O2, Missing}, B::Union{O2, Missing}, A::Union{O3, Missing}, + backend::Bk, allocator::Al + ) where {O1, O2, O3, Bk <: AbstractBackend, Al} + return new{O1, O2, O3, Bk, Al}(D, I, E, C, B, A, backend, allocator) end end -function JordanMPO_AC_Hamiltonian{O1, O2, O3}(D, I, E, C, B, A) where {O1, O2, O3} - return JordanMPO_AC_Hamiltonian{O1, O2, O3}( +function JordanMPO_AC_Hamiltonian{O1, O2, O3}( + D, I, E, C, B, A, backend = DefaultBackend(), allocator = DefaultAllocator() + ) where {O1, O2, O3} + return JordanMPO_AC_Hamiltonian{O1, O2, O3, typeof(backend), typeof(allocator)}( ismissing(D) ? D : convert(O1, D), ismissing(I) ? I : convert(O1, I), ismissing(E) ? E : convert(O1, E), ismissing(C) ? C : convert(O2, C), - ismissing(B) ? E : convert(O2, B), ismissing(A) ? A : convert(O3, A) + ismissing(B) ? B : convert(O2, B), ismissing(A) ? A : convert(O3, A), + backend, allocator ) end function AC_hamiltonian( site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian, above::_HAM_MPS_TYPES, envs; - prepare::Bool = true + prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) @assert below === above "JordanMPO assumptions break" GL = leftenv(envs, site, below) GR = rightenv(envs, site, below) W = operator[site] - H_AC = JordanMPO_AC_Hamiltonian(GL, W, GR) - return prepare ? prepare_operator!!(H_AC) : H_AC + H_AC = JordanMPO_AC_Hamiltonian(GL, W, GR; backend, allocator) + return prepare ? prepare_operator!!(H_AC, backend, allocator) : H_AC end -function JordanMPO_AC_Hamiltonian(GL::MPSTensor, W::JordanMPOTensor, GR::MPSTensor) +function JordanMPO_AC_Hamiltonian( + GL::MPSTensor, W::JordanMPOTensor, GR::MPSTensor; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) # block accessors recompute a fresh `SparseBlockTensorMap` on every access, so bind # them once and reuse the locals throughout WA, WB, WC, WD = W.A, W.B, W.C, W.D @@ -64,7 +74,7 @@ function JordanMPO_AC_Hamiltonian(GL::MPSTensor, W::JordanMPOTensor, GR::MPSTens # starting C = if nonzero_length(WC) > 0 GR_2 = GR[2:(end - 1)] - @plansor starting[-1 -2; -3 -4] ≔ WC[-1; -3 1] * GR_2[-4 1; -2] + @plansor backend = backend allocator = allocator starting[-1 -2; -3 -4] ≔ WC[-1; -3 1] * GR_2[-4 1; -2] only(starting) else missing @@ -73,14 +83,14 @@ function JordanMPO_AC_Hamiltonian(GL::MPSTensor, W::JordanMPOTensor, GR::MPSTens # ending B = if nonzero_length(WB) > 0 GL_2 = GL[2:(end - 1)] - @plansor ending[-1 -2; -3 -4] ≔ GL_2[-1 1; -3] * WB[1 -2; -4] + @plansor backend = backend allocator = allocator ending[-1 -2; -3 -4] ≔ GL_2[-1 1; -3] * WB[1 -2; -4] only(ending) else missing end # continuing - A = MPO_AC_Hamiltonian(GL[2:(end - 1)], WA, GR[2:(end - 1)]) + A = MPO_AC_Hamiltonian(GL[2:(end - 1)], WA, GR[2:(end - 1)], backend, allocator) # obtaining storagetype of environments since these should have already mixed # the types of the operator and state @@ -93,7 +103,7 @@ function JordanMPO_AC_Hamiltonian(GL::MPSTensor, W::JordanMPOTensor, GR::MPSTens # specialization for nearest neighbours nonzero_length(WA) == 0 && (A = missing) - return JordanMPO_AC_Hamiltonian{O1, O2, O3}(D, I, E, C, B, A) + return JordanMPO_AC_Hamiltonian{O1, O2, O3}(D, I, E, C, B, A, backend, allocator) end function prepare_operator!!( @@ -107,11 +117,11 @@ function prepare_operator!!( missing elseif !ismissing(C) Id = TensorKit.id(storagetype(C), space(C, 2)) - @plansor C[-1 -2; -3 -4] += H.D[-1; -3] * Id[-2; -4] + @plansor backend = backend allocator = allocator C[-1 -2; -3 -4] += H.D[-1; -3] * Id[-2; -4] missing elseif !ismissing(B) Id = TensorKit.id(storagetype(B), space(B, 1)) - @plansor B[-1 -2; -3 -4] += Id[-1; -3] * H.D[-2; -4] + @plansor backend = backend allocator = allocator B[-1 -2; -3 -4] += Id[-1; -3] * H.D[-2; -4] missing else H.D @@ -122,7 +132,7 @@ function prepare_operator!!( missing elseif !ismissing(C) Id = id(storagetype(C), space(C, 1)) - @plansor C[-1 -2; -3 -4] += Id[-1; -3] * H.I[-4; -2] + @plansor backend = backend allocator = allocator C[-1 -2; -3 -4] += Id[-1; -3] * H.I[-4; -2] missing else H.I @@ -133,7 +143,7 @@ function prepare_operator!!( missing elseif !ismissing(B) Id = id(storagetype(B), space(B, 2)) - @plansor B[-1 -2; -3 -4] += H.E[-1; -3] * Id[-2; -4] + @plansor backend = backend allocator = allocator B[-1 -2; -3 -4] += H.E[-1; -3] * Id[-2; -4] missing else H.E @@ -142,7 +152,7 @@ function prepare_operator!!( O3′ = prepared_operator_type(O3, typeof(backend), typeof(allocator)) A = ismissing(H.A) ? H.A : prepare_operator!!(H.A, backend, allocator) - return JordanMPO_AC_Hamiltonian{O1, O2, O3′}(D, I, E, C, B, A)::JordanMPO_AC_Hamiltonian{O1, O2, O3′} + return JordanMPO_AC_Hamiltonian{O1, O2, O3′}(D, I, E, C, B, A, backend, allocator)::JordanMPO_AC_Hamiltonian{O1, O2, O3′} end @@ -154,7 +164,7 @@ end Efficient operator for representing the single-site derivative of a `MPOHamiltonian` sandwiched between two MPSs. In particular, this operator aims to make maximal use of the structure of the `MPOHamiltonian` to reduce the number of operations required to apply the operator to a tensor. """ -struct JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4} <: DerivativeOperator +struct JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4, Bk <: AbstractBackend, Al} <: DerivativeOperator II::Union{O1, Missing} # not_started IC::Union{O2, Missing} # starting right ID::Union{O1, Missing} # onsite right @@ -165,38 +175,44 @@ struct JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4} <: DerivativeOperator BE::Union{O2, Missing} # ending left DE::Union{O1, Missing} # onsite left EE::Union{O1, Missing} # finished + backend::Bk # contraction backend used by the matvec + allocator::Al # scratch-buffer allocator used by the matvec - function JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4}( + function JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4, Bk, Al}( II::Union{O1, Missing}, IC::Union{O2, Missing}, ID::Union{O1, Missing}, CB::Union{O2, Missing}, CA::Union{O3, Missing}, AB::Union{O3, Missing}, AA::Union{O4, Missing}, - BE::Union{O2, Missing}, DE::Union{O1, Missing}, EE::Union{O1, Missing} - ) where {O1, O2, O3, O4} - return new{O1, O2, O3, O4}(II, IC, ID, CB, CA, AB, AA, BE, DE, EE) + BE::Union{O2, Missing}, DE::Union{O1, Missing}, EE::Union{O1, Missing}, + backend::Bk, allocator::Al + ) where {O1, O2, O3, O4, Bk <: AbstractBackend, Al} + return new{O1, O2, O3, O4, Bk, Al}(II, IC, ID, CB, CA, AB, AA, BE, DE, EE, backend, allocator) end end function JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4}( - II, IC, ID, CB, CA, AB, AA, BE, DE, EE + II, IC, ID, CB, CA, AB, AA, BE, DE, EE, + backend = DefaultBackend(), allocator = DefaultAllocator() ) where {O1, O2, O3, O4} - return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4}( + return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4, typeof(backend), typeof(allocator)}( ismissing(II) ? II : convert(O1, II), ismissing(IC) ? IC : convert(O2, IC), ismissing(ID) ? ID : convert(O1, ID), ismissing(CB) ? CB : convert(O2, CB), ismissing(CA) ? CA : convert(O3, CA), ismissing(AB) ? AB : convert(O3, AB), ismissing(AA) ? AA : convert(O4, AA), ismissing(BE) ? BE : convert(O2, BE), - ismissing(DE) ? DE : convert(O1, DE), ismissing(EE) ? EE : convert(O1, EE) + ismissing(DE) ? DE : convert(O1, DE), ismissing(EE) ? EE : convert(O1, EE), + backend, allocator ) end function AC2_hamiltonian( site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian, above::_HAM_MPS_TYPES, envs; - prepare::Bool = true + prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) @assert below === above "JordanMPO assumptions break" GL = leftenv(envs, site, below) GR = rightenv(envs, site + 1, below) W1, W2 = operator[site], operator[site + 1] - H_AC2 = JordanMPO_AC2_Hamiltonian(GL, W1, W2, GR) - return prepare ? prepare_operator!!(H_AC2) : H_AC2 + H_AC2 = JordanMPO_AC2_Hamiltonian(GL, W1, W2, GR; backend, allocator) + return prepare ? prepare_operator!!(H_AC2, backend, allocator) : H_AC2 end for f in (:AC_hamiltonian, :AC2_hamiltonian) @@ -208,7 +224,10 @@ for f in (:AC_hamiltonian, :AC2_hamiltonian) end end -function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::JordanMPOTensor, GR::MPSTensor) +function JordanMPO_AC2_Hamiltonian( + GL::MPSTensor, W1::JordanMPOTensor, W2::JordanMPOTensor, GR::MPSTensor; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) # block accessors recompute a fresh `SparseBlockTensorMap` on every access, so bind # them once and reuse the locals throughout A1, B1, C1, D1 = W1.A, W1.B, W1.C, W1.D @@ -222,7 +241,7 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda # starting right IC = if nonzero_length(C2) > 0 - @plansor IC_[-1 -2; -3 -4] ≔ C2[-1; -3 1] * GR[2:(end - 1)][-4 1; -2] + @plansor backend = backend allocator = allocator IC_[-1 -2; -3 -4] ≔ C2[-1; -3 1] * GR[2:(end - 1)][-4 1; -2] only(IC_) else missing @@ -236,7 +255,7 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda # starting left - ending right CB = if nonzero_length(C1) > 0 && nonzero_length(B2) > 0 - @plansor CB_[-1 -2; -3 -4] ≔ C1[-1; -3 1] * B2[1 -2; -4] + @plansor backend = backend allocator = allocator CB_[-1 -2; -3 -4] ≔ C1[-1; -3 1] * B2[1 -2; -4] # have to convert to complex if hamiltonian is real but states are complex scalartype(GL) <: Complex ? complex(only(CB_)) : only(CB_) else @@ -245,7 +264,7 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda # starting left - continuing right CA = if nonzero_length(C1) > 0 && nonzero_length(A2) > 0 - @plansor CA_[-1 -2 -3; -4 -5 -6] ≔ C1[-1; -4 2] * A2[2 -2; -5 1] * + @plansor backend = backend allocator = allocator CA_[-1 -2 -3; -4 -5 -6] ≔ C1[-1; -4 2] * A2[2 -2; -5 1] * GR[2:(end - 1)][-6 1; -3] only(CA_) else @@ -254,7 +273,7 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda # continuing left - ending right AB = if nonzero_length(A1) > 0 && nonzero_length(B2) > 0 - @plansor AB_[-1 -2 -3; -4 -5 -6] ≔ GL[2:(end - 1)][-1 2; -4] * A1[2 -2; -5 1] * + @plansor backend = backend allocator = allocator AB_[-1 -2 -3; -4 -5 -6] ≔ GL[2:(end - 1)][-1 2; -4] * A1[2 -2; -5 1] * B2[1 -3; -6] only(AB_) else @@ -263,14 +282,14 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda # ending left BE = if nonzero_length(B1) > 0 - @plansor BE_[-1 -2; -3 -4] ≔ GL[2:(end - 1)][-1 2; -3] * B1[2 -2; -4] + @plansor backend = backend allocator = allocator BE_[-1 -2; -3 -4] ≔ GL[2:(end - 1)][-1 2; -3] * B1[2 -2; -4] only(BE_) else missing end # continuing - continuing - AA = MPO_AC2_Hamiltonian(GL[2:(end - 1)], A1, A2, GR[2:(end - 1)]) + AA = MPO_AC2_Hamiltonian(GL[2:(end - 1)], A1, A2, GR[2:(end - 1)], backend, allocator) S = spacetype(GL) M = storagetype(GL) @@ -301,7 +320,8 @@ function JordanMPO_AC2_Hamiltonian(GL::MPSTensor, W1::JordanMPOTensor, W2::Jorda II, IC, ID, CB, CA, AB, AA, - BE, DE, EE + BE, DE, EE, + backend, allocator ) end @@ -315,11 +335,11 @@ function prepare_operator!!( CB::Union{Missing, O2} = if !ismissing(CA) && !ismissing(H.CB) Id = TensorKit.id(storagetype(H.CB), space(CA, 3)) - @plansor CA[-1 -2 -3; -4 -5 -6] += H.CB[-1 -2; -4 -5] * Id[-3; -6] + @plansor backend = backend allocator = allocator CA[-1 -2 -3; -4 -5 -6] += H.CB[-1 -2; -4 -5] * Id[-3; -6] missing elseif !ismissing(AB) && !ismissing(H.CB) Id = TensorKit.id(storagetype(H.CB), space(AB, 1)) - @plansor AB[-1 -2 -3; -4 -5 -6] += H.CB[-2 -3; -5 -6] * Id[-1; -4] + @plansor backend = backend allocator = allocator AB[-1 -2 -3; -4 -5 -6] += H.CB[-2 -3; -5 -6] * Id[-1; -4] missing else H.CB @@ -328,7 +348,7 @@ function prepare_operator!!( # starting right IC::Union{Missing, O2} = if !ismissing(CA) && !ismissing(H.IC) Id = TensorKit.id(storagetype(H.IC), space(CA, 1)) - @plansor CA[-1 -2 -3; -4 -5 -6] += Id[-1; -4] * H.IC[ -2 -3; -5 -6] + @plansor backend = backend allocator = allocator CA[-1 -2 -3; -4 -5 -6] += Id[-1; -4] * H.IC[ -2 -3; -5 -6] missing else H.IC @@ -337,7 +357,7 @@ function prepare_operator!!( # ending left BE::Union{Missing, O2} = if !ismissing(AB) && !ismissing(H.BE) Id = TensorKit.id(storagetype(H.BE), space(AB, 3)) - @plansor AB[-1 -2 -3; -4 -5 -6] += H.BE[-1 -2; -4 -5] * Id[-3; -6] + @plansor backend = backend allocator = allocator AB[-1 -2 -3; -4 -5 -6] += H.BE[-1 -2; -4 -5] * Id[-3; -6] missing else H.BE @@ -346,12 +366,12 @@ function prepare_operator!!( # onsite left DE::Union{Missing, O1} = if !ismissing(BE) && !ismissing(H.DE) Id = TensorKit.id(storagetype(H.DE), space(BE, 1)) - @plansor BE[-1 -2; -3 -4] += Id[-1; -3] * H.DE[-2; -4] + @plansor backend = backend allocator = allocator BE[-1 -2; -3 -4] += Id[-1; -3] * H.DE[-2; -4] missing elseif !ismissing(AB) && !ismissing(H.DE) Id1 = id(storagetype(H.DE), space(AB, 1)) Id2 = id(storagetype(H.DE), space(AB, 3)) - @plansor AB[-1 -2 -3; -4 -5 -6] += Id1[-1; -4] * H.DE[-2; -5] * Id2[-3; -6] + @plansor backend = backend allocator = allocator AB[-1 -2 -3; -4 -5 -6] += Id1[-1; -4] * H.DE[-2; -5] * Id2[-3; -6] missing # TODO: could also try in CA? else @@ -361,12 +381,12 @@ function prepare_operator!!( # onsite right ID::Union{Missing, O1} = if !ismissing(IC) && !ismissing(H.ID) Id = TensorKit.id(storagetype(H.ID), space(IC, 2)) - @plansor IC[-1 -2; -3 -4] += H.ID[-1; -3] * Id[-2; -4] + @plansor backend = backend allocator = allocator IC[-1 -2; -3 -4] += H.ID[-1; -3] * Id[-2; -4] missing elseif !ismissing(CA) && !ismissing(H.ID) Id1 = TensorKit.id(storagetype(H.ID), space(CA, 1)) Id2 = TensorKit.id(storagetype(H.ID), space(CA, 3)) - @plansor CA[-1 -2 -3; -4 -5 -6] += Id1[-1; -4] * H.ID[-2; -5] * Id2[-3; -6] + @plansor backend = backend allocator = allocator CA[-1 -2 -3; -4 -5 -6] += Id1[-1; -4] * H.ID[-2; -5] * Id2[-3; -6] missing else H.ID @@ -375,11 +395,11 @@ function prepare_operator!!( # finished II::Union{Missing, O1} = if !ismissing(IC) && !ismissing(H.II) I = id(storagetype(H.II), space(IC, 1)) - @plansor IC[-1 -2; -3 -4] += I[-1; -3] * H.II[-2; -4] + @plansor backend = backend allocator = allocator IC[-1 -2; -3 -4] += I[-1; -3] * H.II[-2; -4] II = missing elseif !ismissing(CA) && !ismissing(H.II) I = id(storagetype(H.II), space(CA, 1) ⊗ space(CA, 2)) - @plansor CA[-1 -2 -3; -4 -5 -6] += I[-1 -2; -4 -5] * H.II[-3; -6] + @plansor backend = backend allocator = allocator CA[-1 -2 -3; -4 -5 -6] += I[-1 -2; -4 -5] * H.II[-3; -6] II = missing else H.II @@ -388,11 +408,11 @@ function prepare_operator!!( # unstarted EE::Union{Missing, O1} = if !ismissing(BE) && !ismissing(H.EE) I = id(storagetype(H.EE), space(BE, 2)) - @plansor BE[-1 -2; -3 -4] += H.EE[-1; -3] * I[-2; -4] + @plansor backend = backend allocator = allocator BE[-1 -2; -3 -4] += H.EE[-1; -3] * I[-2; -4] EE = missing elseif !ismissing(AB) && !ismissing(H.EE) I = id(storagetype(H.EE), space(AB, 2) ⊗ space(AB, 3)) - @plansor AB[-1 -2 -3; -4 -5 -6] += H.EE[-1; -4] * I[-2 -3; -5 -6] + @plansor backend = backend allocator = allocator AB[-1 -2 -3; -4 -5 -6] += H.EE[-1; -4] * I[-2 -3; -5 -6] EE = missing else H.EE @@ -401,35 +421,37 @@ function prepare_operator!!( O4′ = prepared_operator_type(O4, typeof(backend), typeof(allocator)) AA = prepare_operator!!(H.AA, backend, allocator) - return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4′}(II, IC, ID, CB, CA, AB, AA, BE, DE, EE) + return JordanMPO_AC2_Hamiltonian{O1, O2, O3, O4′}(II, IC, ID, CB, CA, AB, AA, BE, DE, EE, backend, allocator) end # Actions # ------- function (H::JordanMPO_AC_Hamiltonian)(x::MPSTensor) + backend, allocator = H.backend, H.allocator y = ismissing(H.A) ? zerovector(x) : H.A(x) - ismissing(H.D) || @plansor y[-1 -2; -3] += x[-1 1; -3] * H.D[-2; 1] - ismissing(H.E) || @plansor y[-1 -2; -3] += H.E[-1; 1] * x[1 -2; -3] - ismissing(H.I) || @plansor y[-1 -2; -3] += x[-1 -2; 1] * H.I[1; -3] - ismissing(H.C) || @plansor y[-1 -2; -3] += x[-1 2; 1] * H.C[-2 -3; 2 1] - ismissing(H.B) || @plansor y[-1 -2; -3] += H.B[-1 -2; 1 2] * x[1 2; -3] + ismissing(H.D) || @plansor backend = backend allocator = allocator y[-1 -2; -3] += x[-1 1; -3] * H.D[-2; 1] + ismissing(H.E) || @plansor backend = backend allocator = allocator y[-1 -2; -3] += H.E[-1; 1] * x[1 -2; -3] + ismissing(H.I) || @plansor backend = backend allocator = allocator y[-1 -2; -3] += x[-1 -2; 1] * H.I[1; -3] + ismissing(H.C) || @plansor backend = backend allocator = allocator y[-1 -2; -3] += x[-1 2; 1] * H.C[-2 -3; 2 1] + ismissing(H.B) || @plansor backend = backend allocator = allocator y[-1 -2; -3] += H.B[-1 -2; 1 2] * x[1 2; -3] return y end function (H::JordanMPO_AC2_Hamiltonian)(x::MPOTensor) + backend, allocator = H.backend, H.allocator y = ismissing(H.AA) ? zerovector(x) : H.AA(x) - ismissing(H.II) || @plansor y[-1 -2; -3 -4] += x[-1 -2; 1 -4] * H.II[-3; 1] - ismissing(H.IC) || @plansor y[-1 -2; -3 -4] += x[-1 -2; 1 2] * H.IC[-4 -3; 2 1] - ismissing(H.ID) || @plansor y[-1 -2; -3 -4] += x[-1 -2; -3 1] * H.ID[-4; 1] - ismissing(H.CB) || @plansor y[-1 -2; -3 -4] += x[-1 1; -3 2] * H.CB[-2 -4; 1 2] - ismissing(H.CA) || @plansor y[-1 -2; -3 -4] += x[-1 1; 3 2] * H.CA[-2 -4 -3; 1 2 3] - ismissing(H.AB) || @plansor y[-1 -2; -3 -4] += x[1 2; -3 3] * H.AB[-1 -2 -4; 1 2 3] - ismissing(H.BE) || @plansor y[-1 -2; -3 -4] += x[1 2; -3 -4] * H.BE[-1 -2; 1 2] - ismissing(H.DE) || @plansor y[-1 -2; -3 -4] += x[-1 1; -3 -4] * H.DE[-2; 1] - ismissing(H.EE) || @plansor y[-1 -2; -3 -4] += x[1 -2; -3 -4] * H.EE[-1; 1] + ismissing(H.II) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 -2; 1 -4] * H.II[-3; 1] + ismissing(H.IC) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 -2; 1 2] * H.IC[-4 -3; 2 1] + ismissing(H.ID) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 -2; -3 1] * H.ID[-4; 1] + ismissing(H.CB) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 1; -3 2] * H.CB[-2 -4; 1 2] + ismissing(H.CA) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 1; 3 2] * H.CA[-2 -4 -3; 1 2 3] + ismissing(H.AB) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[1 2; -3 3] * H.AB[-1 -2 -4; 1 2 3] + ismissing(H.BE) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[1 2; -3 -4] * H.BE[-1 -2; 1 2] + ismissing(H.DE) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[-1 1; -3 -4] * H.DE[-2; 1] + ismissing(H.EE) || @plansor backend = backend allocator = allocator y[-1 -2; -3 -4] += x[1 -2; -3 -4] * H.EE[-1; 1] return y end diff --git a/src/algorithms/derivatives/mpo_derivatives.jl b/src/algorithms/derivatives/mpo_derivatives.jl index 5bc4bc943..c44924e44 100644 --- a/src/algorithms/derivatives/mpo_derivatives.jl +++ b/src/algorithms/derivatives/mpo_derivatives.jl @@ -1,42 +1,68 @@ """ - struct MPODerivativeOperator{L, O <: Tuple, R} + struct MPODerivativeOperator{L, O <: Tuple, R, B, A} Effective local operator obtained from taking the partial derivative of an MPS-MPO-MPS sandwich. + +The `backend` and `allocator` fields are the ones used by the application of the operator. They +default to `DefaultBackend()` and `DefaultAllocator()`, i.e. this operator does not hold on to +any scratch space of its own unless it is explicitly given some. """ -struct MPODerivativeOperator{L, O <: Tuple, R} <: DerivativeOperator +struct MPODerivativeOperator{L, O <: Tuple, R, B <: AbstractBackend, A} <: DerivativeOperator leftenv::L operators::O rightenv::R + backend::B + allocator::A end Base.length(H::MPODerivativeOperator) = length(H.operators) -const MPO_C_Hamiltonian{L, R} = MPODerivativeOperator{L, Tuple{}, R} -MPO_C_Hamiltonian(GL, GR) = MPODerivativeOperator(GL, (), GR) +const MPO_C_Hamiltonian{L, R, B, A} = MPODerivativeOperator{L, Tuple{}, R, B, A} +MPO_C_Hamiltonian( + GL, GR, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = MPODerivativeOperator(GL, (), GR, backend, allocator) -const MPO_AC_Hamiltonian{L, O, R} = MPODerivativeOperator{L, Tuple{O}, R} -MPO_AC_Hamiltonian(GL, O, GR) = MPODerivativeOperator(GL, (O,), GR) +const MPO_AC_Hamiltonian{L, O, R, B, A} = MPODerivativeOperator{L, Tuple{O}, R, B, A} +MPO_AC_Hamiltonian( + GL, O, GR, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = MPODerivativeOperator(GL, (O,), GR, backend, allocator) -const MPO_AC2_Hamiltonian{L, O₁, O₂, R} = MPODerivativeOperator{L, Tuple{O₁, O₂}, R} -MPO_AC2_Hamiltonian(GL, O1, O2, GR) = MPODerivativeOperator(GL, (O1, O2), GR) +const MPO_AC2_Hamiltonian{L, O₁, O₂, R, B, A} = MPODerivativeOperator{L, Tuple{O₁, O₂}, R, B, A} +MPO_AC2_Hamiltonian( + GL, O1, O2, GR, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = MPODerivativeOperator(GL, (O1, O2), GR, backend, allocator) # Constructors # ------------ -function C_hamiltonian(site::Int, below, operator, above, envs; prepare::Bool = true) - H_C = MPO_C_Hamiltonian(leftenv(envs, site + 1, below), rightenv(envs, site, below)) - return prepare ? prepare_operator!!(H_C) : H_C +function C_hamiltonian( + site::Int, below, operator, above, envs; prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + H_C = MPO_C_Hamiltonian( + leftenv(envs, site + 1, below), rightenv(envs, site, below), backend, allocator + ) + return prepare ? prepare_operator!!(H_C, backend, allocator) : H_C end -function AC_hamiltonian(site::Int, below, operator, above, envs; prepare::Bool = true) +function AC_hamiltonian( + site::Int, below, operator, above, envs; prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) O = isnothing(operator) ? nothing : operator[site] - H_AC = MPO_AC_Hamiltonian(leftenv(envs, site, below), O, rightenv(envs, site, below)) - return prepare ? prepare_operator!!(H_AC) : H_AC + H_AC = MPO_AC_Hamiltonian( + leftenv(envs, site, below), O, rightenv(envs, site, below), backend, allocator + ) + return prepare ? prepare_operator!!(H_AC, backend, allocator) : H_AC end -function AC2_hamiltonian(site::Int, below, operator, above, envs; prepare::Bool = true) +function AC2_hamiltonian( + site::Int, below, operator, above, envs; prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) O1, O2 = isnothing(operator) ? (nothing, nothing) : (operator[site], operator[site + 1]) H_AC2 = MPO_AC2_Hamiltonian( - leftenv(envs, site, below), O1, O2, rightenv(envs, site + 1, below) + leftenv(envs, site, below), O1, O2, rightenv(envs, site + 1, below), + backend, allocator ) - return prepare ? prepare_operator!!(H_AC2) : H_AC2 + return prepare ? prepare_operator!!(H_AC2, backend, allocator) : H_AC2 end # Properties @@ -57,55 +83,83 @@ end # Actions # ------- function (h::MPO_C_Hamiltonian{<:MPSBondTensor, <:MPSBondTensor})(x::MPSBondTensor) - @plansor y[-1; -2] ≔ h.leftenv[-1; 1] * x[1; 2] * h.rightenv[2; -2] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1; -2] ≔ h.leftenv[-1; 1] * x[1; 2] * h.rightenv[2; -2] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_C_Hamiltonian{<:MPSTensor, <:MPSTensor})(x::MPSBondTensor) - @plansor y[-1; -2] ≔ h.leftenv[-1 3; 1] * x[1; 2] * h.rightenv[2 3; -2] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1; -2] ≔ h.leftenv[-1 3; 1] * x[1; 2] * h.rightenv[2 3; -2] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor})(x::MPSTensor) - @plansor y[-1 -2; -3] ≔ h.leftenv[-1 5; 4] * x[4 2; 1] * h.operators[1][5 -2; 2 3] * - h.rightenv[1 3; -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3] ≔ h.leftenv[-1 5; 4] * x[4 2; 1] * h.operators[1][5 -2; 2 3] * + h.rightenv[1 3; -3] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC_Hamiltonian{<:MPSTensor, <:Number, <:MPSTensor})(x::MPSTensor) - @plansor y[-1 -2; -3] ≔ ( - h.leftenv[-1 5; 4] * x[4 6; 1] * τ[6 5; 7 -2] * h.rightenv[1 7; -3] - ) * only(h.operators) + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3] ≔ ( + h.leftenv[-1 5; 4] * x[4 6; 1] * τ[6 5; 7 -2] * h.rightenv[1 7; -3] + ) * only(h.operators) + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC_Hamiltonian{<:MPSBondTensor, Nothing, <:MPSBondTensor})(x::MPSTensor) - return @plansor y[-1 -2; -3] ≔ h.leftenv[-1; 2] * x[2 -2; 1] * h.rightenv[1; -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3] ≔ h.leftenv[-1; 2] * x[2 -2; 1] * h.rightenv[1; -3] + end + return y end function (h::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor})( x::GenericMPSTensor{<:Any, 3} ) - @plansor y[-1 -2 -3; -4] ≔ h.leftenv[-1 7; 6] * x[6 4 2; 1] * - h.operators[1][7 -2; 4 5] * τ[5 -3; 2 3] * h.rightenv[1 3; -4] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2 -3; -4] ≔ h.leftenv[-1 7; 6] * x[6 4 2; 1] * + h.operators[1][7 -2; 4 5] * τ[5 -3; 2 3] * h.rightenv[1 3; -4] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC2_Hamiltonian{<:MPSBondTensor, Nothing, Nothing, <:MPSBondTensor})( x::MPOTensor ) - @plansor y[-1 -2; -3 -4] ≔ h.leftenv[-1; 1] * x[1 -2; 2 -4] * h.rightenv[2 -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3 -4] ≔ h.leftenv[-1; 1] * x[1 -2; 2 -4] * h.rightenv[2 -3] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC2_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPOTensor, <:MPSTensor})( x::MPOTensor ) - @plansor y[-1 -2; -3 -4] ≔ h.leftenv[-1 7; 6] * x[6 5; 1 3] * - h.operators[1][7 -2; 5 4] * h.operators[2][4 -4; 3 2] * h.rightenv[1 2; -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3 -4] ≔ h.leftenv[-1 7; 6] * x[6 5; 1 3] * + h.operators[1][7 -2; 5 4] * h.operators[2][4 -4; 3 2] * h.rightenv[1 2; -3] + end return y isa AbstractBlockTensorMap ? only(y) : y end function (h::MPO_AC2_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPOTensor, <:MPSTensor})( x::AbstractTensorMap{<:Any, <:Any, 3, 3} ) - @plansor y[-1 -2 -3; -4 -5 -6] ≔ h.leftenv[-1 11; 10] * x[10 8 6; 1 2 4] * - h.rightenv[1 3; -4] * h.operators[1][11 -2; 8 9] * τ[9 -3; 6 7] * - h.operators[2][7 -6; 4 5] * τ[5 -5; 2 3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + y[-1 -2 -3; -4 -5 -6] ≔ h.leftenv[-1 11; 10] * x[10 8 6; 1 2 4] * + h.rightenv[1 3; -4] * h.operators[1][11 -2; 8 9] * τ[9 -3; 6 7] * + h.operators[2][7 -6; 4 5] * τ[5 -5; 2 3] + end return y isa AbstractBlockTensorMap ? only(y) : y end @@ -135,7 +189,7 @@ VectorInterface.scalartype(::Type{<:PrecomputedDerivative{T}}) where {T} = T TensorKit.storagetype(::Type{<:PrecomputedDerivative{T, S, M}}) where {T, S, M} = M Base.@assume_effects :foldable function prepared_operator_type( - ::Type{MPO_C_Hamiltonian{L, R}}, ::Type{B}, ::Type{A} + ::Type{<:MPO_C_Hamiltonian{L, R}}, ::Type{B}, ::Type{A} ) where {L, R, B, A} T = TensorOperations.promote_contract(scalartype(L), scalartype(R)) S = TensorKit.check_spacetype(L, R) @@ -143,7 +197,7 @@ Base.@assume_effects :foldable function prepared_operator_type( return PrecomputedCDerivative{T, S, M, B, A} end Base.@assume_effects :foldable function prepared_operator_type( - ::Type{MPO_AC_Hamiltonian{L, O, R}}, ::Type{B}, ::Type{A} + ::Type{<:MPO_AC_Hamiltonian{L, O, R}}, ::Type{B}, ::Type{A} ) where {L, O, R, B, A} T = TensorOperations.promote_contract(scalartype(L), scalartype(O), scalartype(R)) S = TensorKit.check_spacetype(L, O, R) @@ -151,7 +205,7 @@ Base.@assume_effects :foldable function prepared_operator_type( return PrecomputedACDerivative{T, S, M, B, A} end Base.@assume_effects :foldable function prepared_operator_type( - ::Type{MPO_AC2_Hamiltonian{L, O₁, O₂, R}}, ::Type{B}, ::Type{A} + ::Type{<:MPO_AC2_Hamiltonian{L, O₁, O₂, R}}, ::Type{B}, ::Type{A} ) where {L, O₁, O₂, R, B, A} T = TensorOperations.promote_contract(scalartype(L), scalartype(O₁), scalartype(O₂), scalartype(R)) S = TensorKit.check_spacetype(L, O₁, O₂, R) diff --git a/src/algorithms/derivatives/projection_derivatives.jl b/src/algorithms/derivatives/projection_derivatives.jl index b3a21fb39..bb5f6945b 100644 --- a/src/algorithms/derivatives/projection_derivatives.jl +++ b/src/algorithms/derivatives/projection_derivatives.jl @@ -1,46 +1,80 @@ -struct ProjectionDerivativeOperator{L, O <: Tuple, R} <: DerivativeOperator +""" + struct ProjectionDerivativeOperator{L, O <: Tuple, R, B, A} + +Effective local operator obtained from taking the partial derivative of the projector `|ψ⟩⟨ψ|` +onto an MPS. + +The `backend` and `allocator` fields are the ones used by the application of the operator. They +default to `DefaultBackend()` and `DefaultAllocator()`, i.e. this operator does not hold on to +any scratch space of its own unless it is explicitly given some. +""" +struct ProjectionDerivativeOperator{L, O <: Tuple, R, B <: AbstractBackend, A} <: DerivativeOperator leftenv::L As::O rightenv::R + backend::B + allocator::A end -const Projection_AC_Hamiltonian{L, O, R} = ProjectionDerivativeOperator{L, Tuple{O}, R} -Projection_AC_Hamiltonian(GL, A, GR) = ProjectionDerivativeOperator(GL, (A,), GR) +const Projection_AC_Hamiltonian{L, O, R, B, A} = ProjectionDerivativeOperator{L, Tuple{O}, R, B, A} +Projection_AC_Hamiltonian( + GL, A, GR, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = ProjectionDerivativeOperator(GL, (A,), GR, backend, allocator) -const Projection_AC2_Hamiltonian{L, O₁, O₂, R} = ProjectionDerivativeOperator{L, Tuple{O₁, O₂}, R} -Projection_AC2_Hamiltonian(GL, A1, A2, GR) = ProjectionDerivativeOperator(GL, (A1, A2), GR) +const Projection_AC2_Hamiltonian{L, O₁, O₂, R, B, A} = + ProjectionDerivativeOperator{L, Tuple{O₁, O₂}, R, B, A} +Projection_AC2_Hamiltonian( + GL, A1, A2, GR, backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = ProjectionDerivativeOperator(GL, (A1, A2), GR, backend, allocator) # Constructors # ------------ function AC_hamiltonian( site::Int, below, operator::ProjectionOperator, above, envs; - prepare::Bool = true + prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) GL = leftenv(envs, site, below) GR = rightenv(envs, site, below) - H_AC = Projection_AC_Hamiltonian(GL, operator.ket.AC[site], GR) - return prepare ? prepare_operator!!(H_AC) : H_AC + H_AC = Projection_AC_Hamiltonian(GL, operator.ket.AC[site], GR, backend, allocator) + return prepare ? prepare_operator!!(H_AC, backend, allocator) : H_AC end -function AC2_hamiltonian(site::Int, below, operator::ProjectionOperator, above, envs; prepare::Bool = true) +function AC2_hamiltonian( + site::Int, below, operator::ProjectionOperator, above, envs; + prepare::Bool = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) GL = leftenv(envs, site, below) GR = rightenv(envs, site + 1, below) - H_AC2 = Projection_AC2_Hamiltonian(GL, operator.ket.AC[site], operator.ket.AR[site + 1], GR) - return prepare ? prepare_operator!!(H_AC2) : H_AC2 + H_AC2 = Projection_AC2_Hamiltonian( + GL, operator.ket.AC[site], operator.ket.AR[site + 1], GR, backend, allocator + ) + return prepare ? prepare_operator!!(H_AC2, backend, allocator) : H_AC2 end # Actions # ------- function (h::Projection_AC_Hamiltonian)(x::MPSTensor) - @plansor v[-1; -2 -3 -4] := h.leftenv[4; -1 -2 5] * h.As[1][5 2; 1] * - h.rightenv[1; -3 -4 3] * conj(x[4 2; 3]) - @plansor y[-1 -2; -3] := conj(v[1; 2 5 6]) * h.leftenv[-1; 1 2 4] * h.As[1][4 -2; 3] * - h.rightenv[3; 5 6 -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + v[-1; -2 -3 -4] := h.leftenv[4; -1 -2 5] * h.As[1][5 2; 1] * + h.rightenv[1; -3 -4 3] * conj(x[4 2; 3]) + end + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3] := conj(v[1; 2 5 6]) * h.leftenv[-1; 1 2 4] * h.As[1][4 -2; 3] * + h.rightenv[3; 5 6 -3] + end return y end function (h::Projection_AC2_Hamiltonian)(x::MPOTensor) - @plansor v[-1; -2 -3 -4] := h.leftenv[6; -1 -2 7] * h.As[1][7 4; 5] * h.As[2][5 2; 1] * - h.rightenv[1; -3 -4 3] * conj(x[6 4; 3 2]) - @plansor y[-1 -2; -3 -4] := conj(v[2; 3 5 6]) * h.leftenv[-1; 2 3 4] * - h.As[1][4 -2; 7] * h.As[2][7 -4; 1] * h.rightenv[1; 5 6 -3] + backend, allocator = h.backend, h.allocator + @plansor backend = backend allocator = allocator begin + v[-1; -2 -3 -4] := h.leftenv[6; -1 -2 7] * h.As[1][7 4; 5] * h.As[2][5 2; 1] * + h.rightenv[1; -3 -4 3] * conj(x[6 4; 3 2]) + end + @plansor backend = backend allocator = allocator begin + y[-1 -2; -3 -4] := conj(v[2; 3 5 6]) * h.leftenv[-1; 2 3 4] * + h.As[1][4 -2; 7] * h.As[2][7 -4; 1] * h.rightenv[1; 5 6 -3] + end return y end diff --git a/src/algorithms/excitation/chepigaansatz.jl b/src/algorithms/excitation/chepigaansatz.jl index 8926aead0..11e57ab43 100644 --- a/src/algorithms/excitation/chepigaansatz.jl +++ b/src/algorithms/excitation/chepigaansatz.jl @@ -24,17 +24,24 @@ Used as the `algorithm` argument of [`excitations`](@ref). * [Chepiga et al. Phys. Rev. B 96 (2017)](@cite chepiga2017) """ -struct ChepigaAnsatz{A <: KrylovAlgorithm} <: Algorithm +struct ChepigaAnsatz{A <: KrylovAlgorithm, B} <: Algorithm "algorithm used for the eigenvalue solvers" alg::A + + "backend for tensor contractions and index manipulations" + backend::B end -function ChepigaAnsatz(; kwargs...) +ChepigaAnsatz(alg::KrylovAlgorithm) = + ChepigaAnsatz(alg, Defaults.backend()) +function ChepigaAnsatz(; + backend = Defaults.backend(), kwargs... + ) if isempty(kwargs) alg = Arnoldi(; krylovdim = 30, tol = 1.0e-10, eager = true) else alg = Arnoldi(; kwargs...) end - return ChepigaAnsatz(alg) + return ChepigaAnsatz(alg, backend) end function excitations( @@ -43,12 +50,13 @@ function excitations( ) 1 ≤ pos ≤ length(ψ) || throw(ArgumentError("invalid position $pos")) isunit(sector) || error("not yet implemented for charged excitations") + allocator = default_allocator(ψ, SerialScheduler()) # add random offset to kickstart Krylov process: AC = ψ.AC[pos] AC₀ = add(AC, randn!(similar(AC)), eps(real(scalartype(AC)))^(1 / 4)) - H_eff = AC_hamiltonian(pos, ψ, H, ψ, envs) + H_eff = AC_hamiltonian(pos, ψ, H, ψ, envs; alg.backend, allocator) Es, ACs, info = eigsolve(H_eff, AC₀, num + 1, :SR, alg.alg) info.converged < num && @warn "excitation failed to converge: normres = $(info.normres)" @@ -93,20 +101,28 @@ Used as the `algorithm` argument of [`excitations`](@ref). * [Chepiga et al. Phys. Rev. B 96 (2017)](@cite chepiga2017) """ -struct ChepigaAnsatz2{A <: KrylovAlgorithm} <: Algorithm +struct ChepigaAnsatz2{A <: KrylovAlgorithm, B} <: Algorithm "algorithm used for the eigenvalue solvers, defaults to `Arnoldi(; krylovdim = 30, tol = 1.0e-10, eager = true)`" alg::A "[truncation strategy](@extref MatrixAlgebraKit.TruncationStrategy) used when splitting the optimized two-site tensor, defaults to `notrunc()`" trunc::Any + + "backend for tensor contractions and index manipulations" + backend::B end -function ChepigaAnsatz2(; trunc = notrunc(), kwargs...) +ChepigaAnsatz2(alg::KrylovAlgorithm, trunc) = + ChepigaAnsatz2(alg, trunc, Defaults.backend()) +function ChepigaAnsatz2(; + trunc = notrunc(), backend = Defaults.backend(), + kwargs... + ) if isempty(kwargs) alg = Arnoldi(; krylovdim = 30, tol = 1.0e-10, eager = true) else alg = Arnoldi(; kwargs...) end - return ChepigaAnsatz2(alg, trunc) + return ChepigaAnsatz2(alg, trunc, backend) end function excitations( @@ -115,12 +131,13 @@ function excitations( ) 1 ≤ pos ≤ length(ψ) - 1 || throw(ArgumentError("invalid position $pos")) isunit(sector) || error("not yet implemented for charged excitations") + allocator = default_allocator(ψ, SerialScheduler()) # add random offset to kickstart Krylov process: @plansor AC2[-1 -2; -3 -4] := ψ.AC[pos][-1 -2; 1] * ψ.AR[pos + 1][1 -4; -3] AC2₀ = add(AC2, randn!(similar(AC2)), eps(real(scalartype(AC2)))^(1 / 4)) - H_eff = AC2_hamiltonian(pos, ψ, H, ψ, envs) + H_eff = AC2_hamiltonian(pos, ψ, H, ψ, envs; alg.backend, allocator) Es, AC2s, info = eigsolve(H_eff, AC2₀, num + 1, :SR, alg.alg) info.converged < num && @warn "excitation failed to converge: normres = $(info.normres)" diff --git a/src/algorithms/excitation/dmrgexcitation.jl b/src/algorithms/excitation/dmrgexcitation.jl index 278585235..521ecbef6 100644 --- a/src/algorithms/excitation/dmrgexcitation.jl +++ b/src/algorithms/excitation/dmrgexcitation.jl @@ -38,6 +38,15 @@ function excitations( states::Tuple{T, Vararg{T}}; init = _perturbed_state(first(states)), num = 1 ) where {T <: FiniteMPS} + # the scope is opened once here rather than inside `_excitations_finite`: the recursion is + # what determines the return type, and a closure at every level would defeat inference + return _excitations_finite(H, alg, states; init, num) +end + +function _excitations_finite( + H::FiniteMPOHamiltonian, alg::FiniteExcited, states::Tuple{T, Vararg{T}}; + init, num + ) where {T <: FiniteMPS} num == 0 && return (scalartype(T)[], T[]) super_op = LinearCombination( @@ -48,7 +57,7 @@ function excitations( ne, _ = find_groundstate(init, super_op, alg.gsalg, envs) nstates = (states..., ne) - ens, excis = excitations(H, alg, nstates; init = init, num = num - 1) + ens, excis = _excitations_finite(H, alg, nstates; init, num = num - 1) pushfirst!(ens, expectation_value(ne, H)) pushfirst!(excis, ne) diff --git a/src/algorithms/grassmann.jl b/src/algorithms/grassmann.jl index add98ffe7..1f55dca38 100644 --- a/src/algorithms/grassmann.jl +++ b/src/algorithms/grassmann.jl @@ -12,7 +12,8 @@ module GrassmannMPS using ..MPSKit using ..MPSKit: AbstractMPSEnvironments, InfiniteEnvironments, MultilineEnvironments, - AC_projection, recalculate!, TimerOutput, DISABLED_TIMER, @timeit + AC_projection, recalculate!, TimerOutput, DISABLED_TIMER, @timeit, default_allocator +using TensorOperations: AbstractBackend, DefaultBackend using TensorKit using OhMyThreads import TensorKitManifolds.Grassmann @@ -137,19 +138,28 @@ function transport!(h, state, g, α::Real, state′) end """ - fg(state, operator, envs = environments(state, operator, state)) + fg(state, operator, envs = environments(state, operator, state); kwargs...) Compute the cost function and the tangent vector with respect to the `AL` parameters of the state. + +For an infinite state the sites are handled concurrently, as decided by the `scheduler` keyword +argument. Note that the scheduler should be passed in from the caller rather than defaulted to, so +that the allocator it selects is inferable. """ function fg( state::FiniteMPS, operator::Union{O, LazySum{O}}, envs::AbstractMPSEnvironments = environments(state, operator, state); timeroutput::TimerOutput = DISABLED_TIMER, + backend::AbstractBackend = DefaultBackend(), + # accepted for interface uniformity, but unused: this gradient is serial regardless + scheduler::Scheduler = SerialScheduler(), ) where {O <: FiniteMPOHamiltonian} f = @timeit timeroutput "expval" expectation_value(state, operator, envs) isapprox(imag(f), 0; atol = eps(abs(f))^(3 / 4)) || @warn "MPO might not be Hermitian: $f" + # the sweep is serial, so a single allocator serves all sites + allocator = default_allocator(state, SerialScheduler()) gs = @timeit timeroutput "gradient" map(1:length(state)) do i - AC′ = AC_projection(i, state, operator, state, envs) + AC′ = AC_projection(i, state, operator, state, envs; backend, allocator) g = Grassmann.project(AC′, state.AL[i]) return rmul(g, state.C[i]') end @@ -159,6 +169,8 @@ function fg( state::InfiniteMPS, operator::Union{O, LazySum{O}}, envs::AbstractMPSEnvironments = environments(state, operator, state); timeroutput::TimerOutput = DISABLED_TIMER, + backend::AbstractBackend = DefaultBackend(), + scheduler::Scheduler = MPSKit.Defaults.scheduler[], ) where {O <: InfiniteMPOHamiltonian} @timeit timeroutput "envs (parallel)" recalculate!(envs, state, operator, state; timeroutput) f = @timeit timeroutput "expval" expectation_value(state, operator, envs) @@ -166,12 +178,12 @@ function fg( A = Core.Compiler.return_type(Grassmann.project, Tuple{eltype(state), eltype(state)}) gs = Vector{A}(undef, length(state)) - @timeit timeroutput "gradient" tmap!( - gs, 1:length(state); scheduler = MPSKit.Defaults.scheduler[] - ) do i - AC′ = AC_projection(i, state, operator, state, envs) + allocator = default_allocator(state, scheduler) + @timeit timeroutput "gradient" tforeach(1:length(state); scheduler) do i + AC′ = AC_projection(i, state, operator, state, envs; backend, allocator) g = Grassmann.project(AC′, state.AL[i]) - return rmul(g, state.C[i]') + gs[i] = rmul(g, state.C[i]') + return nothing end return real(f), gs end @@ -179,6 +191,8 @@ function fg( state::InfiniteMPS, operator::Union{O, LazySum{O}}, envs::AbstractMPSEnvironments = environments(state, operator, state); timeroutput::TimerOutput = DISABLED_TIMER, + backend::AbstractBackend = DefaultBackend(), + scheduler::Scheduler = MPSKit.Defaults.scheduler[], ) where {O <: InfiniteMPO} @timeit timeroutput "envs (parallel)" recalculate!(envs, state, operator, state; timeroutput) f = @timeit timeroutput "expval" expectation_value(state, operator, envs) @@ -186,12 +200,12 @@ function fg( A = Core.Compiler.return_type(Grassmann.project, Tuple{eltype(state), eltype(state)}) gs = Vector{A}(undef, length(state)) - @timeit timeroutput "gradient" tmap!( - gs, eachindex(state); scheduler = MPSKit.Defaults.scheduler[] - ) do i - AC′ = AC_projection(i, state, operator, state, envs) + allocator = default_allocator(state, scheduler) + @timeit timeroutput "gradient" tforeach(eachindex(state); scheduler) do i + AC′ = AC_projection(i, state, operator, state, envs; backend, allocator) g = rmul!(Grassmann.project(AC′, state.AL[i]), -inv(f)) - return rmul(g, state.C[i]') + gs[i] = rmul(g, state.C[i]') + return nothing end return -log(real(f)), gs end @@ -199,6 +213,8 @@ function fg( state::MultilineMPS, operator::MultilineMPO, envs::MultilineEnvironments = environments(state, operator, state); timeroutput::TimerOutput = DISABLED_TIMER, + backend::AbstractBackend = DefaultBackend(), + scheduler::Scheduler = MPSKit.Defaults.scheduler[], ) @assert length(state) == 1 "not implemented" @timeit timeroutput "envs (parallel)" recalculate!(envs, state, operator, state; timeroutput) @@ -207,10 +223,9 @@ function fg( A = Core.Compiler.return_type(Grassmann.project, Tuple{eltype(state), eltype(state)}) gs = Matrix{A}(undef, size(state)) - @timeit timeroutput "gradient" tforeach( - eachindex(state); scheduler = MPSKit.Defaults.scheduler[] - ) do i - AC′ = AC_projection(i, state, operator, state, envs) + allocator = default_allocator(state, scheduler) + @timeit timeroutput "gradient" tforeach(eachindex(state); scheduler) do i + AC′ = AC_projection(i, state, operator, state, envs; backend, allocator) g = rmul!(Grassmann.project(AC′, state.AL[i]), -inv(f)) gs[i] = rmul(g, state.C[i]') return nothing diff --git a/src/algorithms/groundstate/dmrg.jl b/src/algorithms/groundstate/dmrg.jl index 034822208..9c996dbcb 100644 --- a/src/algorithms/groundstate/dmrg.jl +++ b/src/algorithms/groundstate/dmrg.jl @@ -59,7 +59,7 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`find_groundstate`](@ref) and [`approximate`](@ref). """ -struct DMRG{A, F, E, G} <: Algorithm +struct DMRG{A, F, E, G, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 @@ -82,12 +82,16 @@ struct DMRG{A, F, E, G} <: Algorithm algorithm with no truncation, or a truncated SVD), or an algorithm that additionally expands the bond beforehand (e.g. [`DMRG3S`](@ref))" alg_gauge::G + + "backend for tensor contractions and index manipulations" + backend::B end function DMRG(; tol = Defaults.tol, maxiter = Defaults.maxiter, alg_eigsolve = (;), verbosity = Defaults.verbosity, finalize = Defaults._finalize, alg_expand = nothing, alg_gauge = nothing, trunc = nothing, - alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth() + alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth(), + backend = Defaults.backend() ) # single-site DMRG defaults to the per-bond adaptive controller (`AdaptiveKrylov`); pass # `alg_eigsolve = (; adaptive = false, ...)` to opt out (the splat overrides the default). @@ -111,34 +115,39 @@ function DMRG(; if (!isnothing(alg_expand) || _expands(alg_gauge)) && !_truncates(alg_gauge) @warn "DMRG with a bond-expanding `alg_expand` and/or `alg_gauge` but no truncation (`trunc = notrunc()`): the bond dimension will grow unboundedly each sweep." end - return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize, alg_expand, alg_gauge) + return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize, alg_expand, alg_gauge, backend) end - function local_update!( site, direction, ψ, O, alg::DMRG, envs, ϵ_global, ϵ_trunc, decay_rate, - iter, timeroutput + iter, timeroutput, allocator ) - ϵ_local = calc_galerkin(site, ψ, O, ψ, envs) + ϵ_local = calc_galerkin(site, ψ, O, ψ, envs; alg.backend, allocator) # 1. expand - isnothing(alg.alg_expand) || - @timeit timeroutput "expand" changebond!(site, direction, ψ, O, alg.alg_expand, envs) + if !isnothing(alg.alg_expand) + @timeit timeroutput "expand" changebond!( + site, direction, ψ, O, alg.alg_expand, envs; allocator + ) + end # 2. local update alg_eigsolve = adapt_solver(alg.alg_eigsolve; decay_rate, g_local = ϵ_local, g_global = ϵ_global, eps_trunc = ϵ_trunc) ac_old = ψ.AC[site] λ, AC′, info = @timeit timeroutput "AC_eigsolve" begin - H_effective = AC_hamiltonian(site, ψ, O, ψ, envs) + H_effective = AC_hamiltonian(site, ψ, O, ψ, envs; alg.backend, allocator) fixedpoint(H_effective, ac_old, :SR, alg_eigsolve) end alg_gauge = _update_alg_gauge(alg.alg_gauge, iter, ϵ_global) # 3. gauge - ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge!(ψ, site, direction, O, envs, AC′, alg_gauge; normalize = true) + ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge!( + ψ, site, direction, O, envs, AC′, alg_gauge; + normalize = true, alg.backend, allocator + ) # 4. bookkeeping: measured contraction factor per matvec, kept a strict contraction in (0, 1) decay_rate = clamp((first(info.normres) / ϵ_local)^(1 / max(1, info.numops)), 1.0e-3, 0.999) @@ -161,7 +170,7 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`find_groundstate`](@ref) and [`approximate`](@ref). """ -struct DMRG2{A, G, F} <: Algorithm +struct DMRG2{A, G, F, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 @@ -179,12 +188,16 @@ struct DMRG2{A, G, F} <: Algorithm "callback function applied after each iteration, of signature `finalize(iter, ψ, H, envs) -> ψ, envs`" finalize::F + + "backend for tensor contractions and index manipulations" + backend::B end # TODO: find better default truncation function DMRG2(; tol = Defaults.tol, maxiter = Defaults.maxiter, verbosity = Defaults.verbosity, alg_eigsolve = (;), alg_svd = Defaults.alg_svd(), trunc, - finalize = Defaults._finalize + finalize = Defaults._finalize, + backend = Defaults.backend() ) # two-site DMRG defaults to the per-bond adaptive controller (`AdaptiveKrylov`); pass # `alg_eigsolve = (; adaptive = false, ...)` to opt out (the splat overrides the default). @@ -192,16 +205,16 @@ function DMRG2(; Defaults.alg_eigsolve(; adaptive = true, alg_eigsolve...) : alg_eigsolve # two-site DMRG always truncates the enlarged bond back down, so the gauge is a truncated SVD alg_gauge = MatrixAlgebraKit.TruncatedAlgorithm(alg_svd, trunc) - return DMRG2(tol, maxiter, verbosity, alg_eigsolve′, alg_gauge, finalize) + return DMRG2(tol, maxiter, verbosity, alg_eigsolve′, alg_gauge, finalize, backend) end function local_update!( pos, direction, ψ, O, alg::DMRG2, envs, ϵ_global, ϵ_trunc, decay_rate, - iter, timeroutput + iter, timeroutput, allocator ) - Heff = @timeit timeroutput "AC2_hamiltonian" AC2_hamiltonian(pos, ψ, O, ψ, envs) + Heff = @timeit timeroutput "AC2_hamiltonian" AC2_hamiltonian(pos, ψ, O, ψ, envs; alg.backend, allocator) kind = direction === Val(:right) ? :ACAR : :ALAC ac2 = AC2(ψ, pos; kind) @@ -265,6 +278,14 @@ Currently supported for the finite-system algorithms [`DMRG`](@ref) and [`DMRG2` function find_groundstate!( ψ::AbstractFiniteMPS, H, alg::Union{DMRG, DMRG2}, envs = environments(ψ, H, ψ) ) + # the sweep is serial, so a single allocator serves all local updates + allocator = default_allocator(ψ, SerialScheduler()) + return _find_groundstate_sweep!(ψ, H, alg, envs, allocator) +end + +function _find_groundstate_sweep!( + ψ::AbstractFiniteMPS, H, alg::Union{DMRG, DMRG2}, envs, allocator + ) name = string(nameof(typeof(alg))) log = IterLog(name) timeroutput = TimerOutput(name) @@ -289,7 +310,7 @@ function find_groundstate!( pos, Val(:right), ψ, H, alg, envs, ϵ_global, ϵ_truncs[pos], decay_rates[pos], - iter, timeroutput + iter, timeroutput, allocator ) ϵ_global = maximum(ϵ_locals) end @@ -301,7 +322,7 @@ function find_groundstate!( pos, Val(:left), ψ, H, alg, envs, ϵ_global, ϵ_truncs[pos], decay_rates[pos], - iter, timeroutput + iter, timeroutput, allocator ) ϵ_global = maximum(ϵ_locals) end diff --git a/src/algorithms/groundstate/gradient_grassmann.jl b/src/algorithms/groundstate/gradient_grassmann.jl index c5be94d82..e295df43e 100644 --- a/src/algorithms/groundstate/gradient_grassmann.jl +++ b/src/algorithms/groundstate/gradient_grassmann.jl @@ -17,6 +17,7 @@ with a preconditioner to induce the metric from the Hilbert space inner product. - `tol = Defaults.tol`: tolerance for convergence criterium - `maxiter = Defaults.maxiter`: maximum amount of iterations - `verbosity = Defaults.verbosity - 1`: level of information display +- `backend = Defaults.backend()`: backend for tensor contractions and index manipulations # Fields @@ -30,16 +31,19 @@ Used as the `algorithm` argument of [`find_groundstate`](@ref) and [`leading_bou * [Hauru et al. SciPost Phys. 10 (2021)](@cite hauru2021) """ -struct GradientGrassmann{O <: OptimKit.OptimizationAlgorithm, F} <: Algorithm +struct GradientGrassmann{O <: OptimKit.OptimizationAlgorithm, F, B} <: Algorithm "optimization algorithm" method::O "callback function applied after each iteration, of signature `finalize!(x, f, g, numiter) -> x, f, g`" finalize!::F + "backend for tensor contractions and index manipulations" + backend::B function GradientGrassmann(; method = ConjugateGradient, (finalize!) = OptimKit._finalize!, tol = Defaults.tol, maxiter = Defaults.maxiter, - verbosity = Defaults.verbosity - 1 + verbosity = Defaults.verbosity - 1, + backend = Defaults.backend() ) if isa(method, OptimKit.OptimizationAlgorithm) # We were given an optimisation method, just use it. @@ -55,7 +59,7 @@ struct GradientGrassmann{O <: OptimKit.OptimizationAlgorithm, F} <: Algorithm msg = "method should be either an instance or a subtype of `OptimKit.OptimizationAlgorithm`." throw(ArgumentError(msg)) end - return new{typeof(m), typeof(finalize!)}(m, finalize!) + return new{typeof(m), typeof(finalize!), typeof(backend)}(m, finalize!, backend) end end @@ -70,7 +74,12 @@ function find_groundstate( method_verbosity = hasproperty(alg.method, :verbosity) ? alg.method.verbosity : 0 method_verbosity > 3 || disable_timer!(timeroutput) - fg(x) = timeit(() -> GrassmannMPS.fg(x, H, envs; timeroutput), timeroutput, "fg") + # read the scheduler here rather than in `fg`, so that the allocator it selects is inferable + scheduler = Defaults.scheduler[] + fg(x) = timeit( + () -> GrassmannMPS.fg(x, H, envs; timeroutput, alg.backend, scheduler), + timeroutput, "fg", + ) retract(state, g, α) = timeit( () -> GrassmannMPS.retract(state, g, α), timeroutput, "retract", ) diff --git a/src/algorithms/groundstate/idmrg.jl b/src/algorithms/groundstate/idmrg.jl index a7f8d326a..f44770a1f 100644 --- a/src/algorithms/groundstate/idmrg.jl +++ b/src/algorithms/groundstate/idmrg.jl @@ -11,7 +11,7 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`find_groundstate`](@ref), [`leading_boundary`](@ref), and [`approximate`](@ref). """ -@kwdef struct IDMRG{A} <: Algorithm +@kwdef struct IDMRG{A, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 = Defaults.tol @@ -26,6 +26,9 @@ Used as the `algorithm` argument of [`find_groundstate`](@ref), [`leading_bounda "algorithm used for the eigenvalue solvers" alg_eigsolve::A = Defaults.alg_eigsolve() + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end """ @@ -41,7 +44,7 @@ $(TYPEDFIELDS) Used as the `algorithm` argument of [`find_groundstate`](@ref), [`leading_boundary`](@ref), and [`approximate`](@ref). """ -@kwdef struct IDMRG2{A, S} <: Algorithm +@kwdef struct IDMRG2{A, S, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 = Defaults.tol @@ -62,11 +65,13 @@ Used as the `algorithm` argument of [`find_groundstate`](@ref), [`leading_bounda "algorithm used for [truncation](@extref MatrixAlgebraKit.TruncationStrategy) of the two-site update" trunc::TruncationStrategy -end + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() +end # Internal state of the IDMRG algorithm -struct IDMRGState{S, O, E, T} +struct IDMRGState{S, O, E, T, A} mps::S operator::O envs::E @@ -74,15 +79,25 @@ struct IDMRGState{S, O, E, T} ϵ::Float64 # TODO: Could be any <:Real energy::T timeroutput::TimerOutput + # the sweeps are serial, so a single allocator serves every local update of every iteration: + # it is obtained once when the solve starts and carried along here, as reusing one buffer + # rather than growing a fresh one each sweep is the entire point of having it + allocator::A end function IDMRGState{T}( mps::S, operator::O, envs::E, iter::Int, ϵ::Float64, energy, - timeroutput::TimerOutput, - ) where {S, O, E, T} - return IDMRGState{S, O, E, T}(mps, operator, envs, iter, ϵ, T(energy), timeroutput) + timeroutput::TimerOutput, allocator::A, + ) where {S, O, E, T, A} + return IDMRGState{S, O, E, T, A}( + mps, operator, envs, iter, ϵ, T(energy), timeroutput, allocator + ) end function find_groundstate(mps, operator, alg::alg_type, envs = environments(mps, operator, mps)) where {alg_type <: Union{<:IDMRG, <:IDMRG2}} + return _find_groundstate_idmrg(mps, operator, alg, envs) +end + +function _find_groundstate_idmrg(mps, operator, alg::alg_type, envs) where {alg_type <: Union{<:IDMRG, <:IDMRG2}} (length(mps) ≤ 1 && alg isa IDMRG2) && throw(ArgumentError("unit cell should be >= 2")) name = alg isa IDMRG ? "IDMRG" : "IDMRG2" log = IterLog(name) @@ -90,7 +105,9 @@ function find_groundstate(mps, operator, alg::alg_type, envs = environments(mps, alg.verbosity > 3 || disable_timer!(timeroutput) mps = copy(mps) iter = 0 - ϵ = calc_galerkin(mps, operator, mps, envs) + # the sweeps are serial, so one allocator serves the whole solve + allocator = default_allocator(mps, SerialScheduler()) + ϵ = calc_galerkin(mps, operator, mps, envs; alg.backend, allocator) E = zero(TensorOperations.promote_contract(scalartype(mps), scalartype(operator))) LoggingExtras.withlevel(; alg.verbosity) do @@ -100,7 +117,7 @@ function find_groundstate(mps, operator, alg::alg_type, envs = environments(mps, end end - state = IDMRGState(mps, operator, envs, iter, ϵ, E, timeroutput) + state = IDMRGState(mps, operator, envs, iter, ϵ, E, timeroutput, allocator) it = IterativeSolver(alg, state) return LoggingExtras.withlevel(; alg.verbosity) do @@ -150,7 +167,7 @@ function Base.iterate( # update state it.state = IDMRGState{T}( - mps, state.operator, envs, state.iter + 1, ϵ, E_new, timeroutput, + mps, state.operator, envs, state.iter + 1, ϵ, E_new, timeroutput, state.allocator, ) return (mps, envs, ϵ, ΔE), it.state @@ -161,7 +178,8 @@ function localupdate_step!( ) alg_eigsolve = adapt_solver(it.alg_eigsolve; iter = state.iter, g_global = state.ϵ) return _localupdate_sweep_idmrg!( - state.mps, state.operator, state.envs, alg_eigsolve, state.timeroutput, + state.mps, state.operator, state.envs, alg_eigsolve, state.timeroutput; + it.backend, state.allocator, ) end @@ -171,17 +189,21 @@ function localupdate_step!( alg_eigsolve = adapt_solver(it.alg_eigsolve; iter = state.iter, g_global = state.ϵ) return _localupdate_sweep_idmrg2!( state.mps, state.operator, state.envs, alg_eigsolve, - it.trunc, it.alg_svd, state.timeroutput, + it.trunc, it.alg_svd, state.timeroutput; + it.backend, state.allocator, ) end -function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, timeroutput::TimerOutput) +function _localupdate_sweep_idmrg!( + ψ, H, envs, alg_eigsolve, timeroutput::TimerOutput; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) local E C_old = ψ.C[0] # left to right sweep for pos in 1:length(ψ) @timeit timeroutput "AC_eigsolve" begin - h = AC_hamiltonian(pos, ψ, H, ψ, envs) + h = AC_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator) _, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve) end @timeit timeroutput "ortho_step" begin @@ -198,7 +220,7 @@ function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, timeroutput::Timer # right to left sweep for pos in length(ψ):-1:1 @timeit timeroutput "AC_eigsolve" begin - h = AC_hamiltonian(pos, ψ, H, ψ, envs) + h = AC_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator) E, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve) end @timeit timeroutput "ortho_step" begin @@ -210,9 +232,9 @@ function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, timeroutput::Timer return ψ, envs, C_old, E end - function _localupdate_sweep_idmrg2!( - ψ, H, envs, alg_eigsolve, alg_trunc, alg_svd, timeroutput::TimerOutput, + ψ, H, envs, alg_eigsolve, alg_trunc, alg_svd, timeroutput::TimerOutput; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) # @timeit wraps its body in try-finally, which is a new lexical scope: declare locals # at function scope so values can flow between consecutive @timeit blocks. @@ -221,7 +243,7 @@ function _localupdate_sweep_idmrg2!( for pos in 1:(length(ψ) - 1) @timeit timeroutput "AC2_eigsolve" begin ac2 = AC2(ψ, pos; kind = :ACAR) - h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs) + h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator) _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve) end @timeit timeroutput "svd_trunc" begin @@ -244,7 +266,7 @@ function _localupdate_sweep_idmrg2!( ψ.AC[1] = _mul_tail(ψ.AL[1], ψ.C[1]) @timeit timeroutput "AC2_eigsolve" begin ac2 = AC2(ψ, 0; kind = :ALAC) - h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs) + h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs; backend, allocator) _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve) end @timeit timeroutput "svd_trunc" begin @@ -272,7 +294,7 @@ function _localupdate_sweep_idmrg2!( for pos in (length(ψ) - 1):-1:1 @timeit timeroutput "AC2_eigsolve" begin ac2 = AC2(ψ, pos; kind = :ALAC) - h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs) + h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator) _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve) end @timeit timeroutput "svd_trunc" begin @@ -296,7 +318,7 @@ function _localupdate_sweep_idmrg2!( ψ.AR[1] = _transpose_front(ψ.C[end] \ _transpose_tail(ψ.AC[1])) @timeit timeroutput "AC2_eigsolve" begin ac2 = AC2(ψ, 0; kind = :ACAR) - h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs) + h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs; backend, allocator) E, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve) end @timeit timeroutput "svd_trunc" begin diff --git a/src/algorithms/groundstate/vumps.jl b/src/algorithms/groundstate/vumps.jl index 149194c06..716d2e21d 100644 --- a/src/algorithms/groundstate/vumps.jl +++ b/src/algorithms/groundstate/vumps.jl @@ -16,7 +16,7 @@ Used as the `algorithm` argument of [`find_groundstate`](@ref) and [`leading_bou * [Zauner-Stauber et al. Phys. Rev. B 97 (2018)](@cite zauner-stauber2018) * [Vanderstraeten et al. SciPost Phys. Lect. Notes 7 (2019)](@cite vanderstraeten2019) """ -@kwdef struct VUMPS{F} <: Algorithm +@kwdef struct VUMPS{F, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 = Defaults.tol @@ -37,6 +37,9 @@ Used as the `algorithm` argument of [`find_groundstate`](@ref) and [`leading_bou "callback function applied after each iteration, of signature `finalize(iter, ψ, H, envs) -> ψ, envs`" finalize::F = Defaults._finalize + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end # Internal state of the VUMPS algorithm @@ -66,7 +69,7 @@ function dominant_eigsolve( iter = 0 mps = copy(mps) - ϵ = calc_galerkin(mps, operator, mps, envs) + ϵ = calc_galerkin(mps, operator, mps, envs; alg.backend) alg_environments = adapt_solver(alg.alg_environments; iter, g_global = ϵ) recalculate!(envs, mps, operator, mps, alg_environments; timeroutput) @@ -109,7 +112,9 @@ function Base.iterate(it::IterativeSolver{<:VUMPS}, state = it.state) )::typeof((mps, envs)) # error criterion - ϵ = @timeit timeroutput "calc_galerkin" calc_galerkin(mps, state.operator, mps, envs) + ϵ = @timeit timeroutput "calc_galerkin" calc_galerkin( + mps, state.operator, mps, envs; it.backend + ) # update state it.state = VUMPSState( @@ -132,14 +137,14 @@ function localupdate_step!( ACs = mps.AL dst_ACs = mps isa Multiline ? eachcol(ACs) : ACs - tree_point = String[section.name for section in state.timeroutput.timer_stack] - tforeach(eachsite(mps), src_ACs, src_Cs; scheduler) do site, AC₀, C₀ + allocator = default_allocator(mps, scheduler) + tforeach(eachsite(mps); scheduler) do site sub_timeroutput = TimerOutput() dst_ACs[site] = _localupdate_vumps_step!( - site, mps, state.operator, state.envs, AC₀, C₀; - parallel = false, alg_orth, state.which, alg_eigsolve, - timeroutput = sub_timeroutput, + site, mps, state.operator, state.envs, src_ACs[site], src_Cs[site]; + alg_orth, state.which, alg_eigsolve, + timeroutput = sub_timeroutput, it.backend, allocator, ) state.timeroutput.enabled && merge!(state.timeroutput, sub_timeroutput; tree_point) @@ -150,41 +155,19 @@ end function _localupdate_vumps_step!( site, mps, operator, envs, AC₀, C₀; - parallel::Bool = false, alg_orth = Defaults.alg_orth(), + alg_orth = Defaults.alg_orth(), alg_eigsolve = Defaults.eigsolver, which, timeroutput::TimerOutput = DISABLED_TIMER, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator(), ) - if !parallel - local AC, C - @timeit timeroutput "AC_eigsolve" begin - Hac = AC_hamiltonian(site, mps, operator, mps, envs) - _, AC = fixedpoint(Hac, AC₀, which, alg_eigsolve) - end - @timeit timeroutput "C_eigsolve" begin - Hc = C_hamiltonian(site, mps, operator, mps, envs) - _, C = fixedpoint(Hc, C₀, which, alg_eigsolve) - end - return regauge!(AC, C; alg = alg_orth) - end - local AC, C - @sync begin - @spawn begin - sub_timeroutput = TimerOutput() - @timeit sub_timeroutput "AC_eigsolve" begin - Hac = AC_hamiltonian(site, mps, operator, mps, envs) - _, AC = fixedpoint(Hac, AC₀, which, alg_eigsolve) - end - timeroutput.enabled && merge!(timeroutput, sub_timeroutput) - end - @spawn begin - sub_timeroutput = TimerOutput() - @timeit sub_timeroutput "C_eigsolve" begin - Hc = C_hamiltonian(site, mps, operator, mps, envs) - _, C = fixedpoint(Hc, C₀, which, alg_eigsolve) - end - timeroutput.enabled && merge!(timeroutput, sub_timeroutput) - end + @timeit timeroutput "AC_eigsolve" begin + Hac = AC_hamiltonian(site, mps, operator, mps, envs; backend, allocator) + _, AC = fixedpoint(Hac, AC₀, which, alg_eigsolve) + end + @timeit timeroutput "C_eigsolve" begin + Hc = C_hamiltonian(site, mps, operator, mps, envs; backend, allocator) + _, C = fixedpoint(Hc, C₀, which, alg_eigsolve) end return regauge!(AC, C; alg = alg_orth) end diff --git a/src/algorithms/post_expand/dmrg3s.jl b/src/algorithms/post_expand/dmrg3s.jl index 065c8f16e..6ab38c238 100644 --- a/src/algorithms/post_expand/dmrg3s.jl +++ b/src/algorithms/post_expand/dmrg3s.jl @@ -130,7 +130,11 @@ function _get_combiner(::Type{TorA}, V1, V2) where {TorA} return isomorphism(TorA, oplus(fuse(Vprod)), Vprod) end -function gauge!(ψ::AbstractFiniteMPS, pos::Int, ::Val{:right}, H, envs, AC, alg::DMRG3S; normalize = true) +function gauge!( + ψ::AbstractFiniteMPS, pos::Int, ::Val{:right}, H, envs, AC, alg::DMRG3S; + normalize = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) El = leftenv(envs, pos, ψ) Hi = H[pos] α = alg.noise @@ -140,7 +144,7 @@ function gauge!(ψ::AbstractFiniteMPS, pos::Int, ::Val{:right}, H, envs, AC, alg combiner = _get_combiner(T, V, right_virtualspace(Hi))' Vpert = only(domain(combiner)) - mpo_ac = MPO_AC_Hamiltonian(El, Hi, combiner) + mpo_ac = MPO_AC_Hamiltonian(El, Hi, combiner, backend, allocator) pert = α * mpo_ac(AC) AC_expanded = catdomain(AC, pert) @@ -155,7 +159,11 @@ function gauge!(ψ::AbstractFiniteMPS, pos::Int, ::Val{:right}, H, envs, AC, alg return ψ, ϵ end -function gauge!(ψ::AbstractFiniteMPS, pos::Int, ::Val{:left}, H, envs, AC, alg::DMRG3S; normalize = true) +function gauge!( + ψ::AbstractFiniteMPS, pos::Int, ::Val{:left}, H, envs, AC, alg::DMRG3S; + normalize = true, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) Er = rightenv(envs, pos, ψ) Hi = H[pos] α = alg.noise @@ -166,7 +174,7 @@ function gauge!(ψ::AbstractFiniteMPS, pos::Int, ::Val{:left}, H, envs, AC, alg: Vpert = only(codomain(combiner)) combiner = _transpose_front(combiner) - mpo_ac = MPO_AC_Hamiltonian(combiner, Hi, Er) + mpo_ac = MPO_AC_Hamiltonian(combiner, Hi, Er, backend, allocator) pert = α * mpo_ac(AC) AC = _transpose_tail(AC) pert = _transpose_tail(pert) diff --git a/src/algorithms/post_expand/post_expand.jl b/src/algorithms/post_expand/post_expand.jl index 5f8579c37..e15954d25 100644 --- a/src/algorithms/post_expand/post_expand.jl +++ b/src/algorithms/post_expand/post_expand.jl @@ -3,7 +3,13 @@ _update_alg_gauge(alg, iter, ϵ) = alg set_alg_gauge(::Nothing, inner_gauge) = inner_gauge alg_gauge(alg) = alg -gauge!(ψ::AbstractFiniteMPS, pos::Int, direction, H, envs, AC, alg_gauge; normalize::Bool = false) = - gauge!(ψ, pos, direction, AC, alg_gauge; normalize) +# A plain (non-expanding) gauge step does not touch `H`/`envs`, and thus has no contraction to +# route a `backend`/`allocator` to: accept and drop them, so that callers can pass them +# unconditionally and only the expanding gauges (e.g. [`DMRG3S`](@ref)) pick them up. +gauge!( + ψ::AbstractFiniteMPS, pos::Int, direction, H, envs, AC, alg_gauge; + normalize::Bool = false, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() +) = gauge!(ψ, pos, direction, AC, alg_gauge; normalize) gauge2!(ψ::AbstractFiniteMPS, pos::Int, direction, H, envs, AC2, alg_gauge; normalize::Bool = false) = gauge2!(ψ, pos, direction, AC2, alg_gauge; normalize) diff --git a/src/algorithms/propagator/corvector.jl b/src/algorithms/propagator/corvector.jl index 89466e7e0..6b4f3066b 100644 --- a/src/algorithms/propagator/corvector.jl +++ b/src/algorithms/propagator/corvector.jl @@ -23,7 +23,7 @@ Used as the `algorithm` argument of [`propagator`](@ref). * [Jeckelmann. Phys. Rev. B 66 (2002)](@cite jeckelmann2002) """ -@kwdef struct DynamicalDMRG{F <: DDMRG_Flavour, S} <: Algorithm +@kwdef struct DynamicalDMRG{F <: DDMRG_Flavour, S, B} <: Algorithm "flavour of the algorithm to use, either of type [`NaiveInvert`](@ref) or [`Jeckelmann`](@ref)" flavour::F = NaiveInvert() "algorithm used for the linear solvers" @@ -34,6 +34,8 @@ Used as the `algorithm` argument of [`propagator`](@ref). maxiter::Int = Defaults.maxiter "setting for how much information is displayed" verbosity::Int = Defaults.verbosity + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end """ @@ -71,6 +73,7 @@ function propagator( A::AbstractFiniteMPS, z::Number, H, alg::DynamicalDMRG{NaiveInvert}; init = copy(A) ) + allocator = default_allocator(A, SerialScheduler()) h_envs = environments(init, H, init) # environments for h mixedenvs = environments(init, A) # environments for @@ -83,9 +86,9 @@ function propagator( ϵ = 0.0 for i in [1:(length(A) - 1); length(A):-1:2] - tos = AC_projection(i, init, A, mixedenvs) + tos = AC_projection(i, init, A, mixedenvs; alg.backend, allocator) - H_AC = AC_hamiltonian(i, init, H, init, h_envs) + H_AC = AC_hamiltonian(i, init, H, init, h_envs; alg.backend, allocator) AC = init.AC[i] AC′, convhist = linsolve(H_AC, -tos, AC, alg.solver, -z, one(z)) @@ -143,6 +146,7 @@ function propagator( A::AbstractFiniteMPS, z::Number, H, alg::DynamicalDMRG{Jeckelmann}; init = copy(A) ) + allocator = default_allocator(A, SerialScheduler()) ω = real(z) η = imag(z) @@ -159,9 +163,9 @@ function propagator( ϵ = 0.0 for i in [1:(length(A) - 1); length(A):-1:2] - tos = AC_projection(i, init, A, mixedenvs) - H1_AC = AC_hamiltonian(i, init, H, init, envs1) - H2_AC = AC_hamiltonian(i, init, H2, init, envs2) + tos = AC_projection(i, init, A, mixedenvs; alg.backend, allocator) + H1_AC = AC_hamiltonian(i, init, H, init, envs1; alg.backend, allocator) + H2_AC = AC_hamiltonian(i, init, H2, init, envs2; alg.backend, allocator) H_AC = LinearCombination((H1_AC, H2_AC), (-2 * ω, 1)) AC′, convhist = linsolve(H_AC, -η * tos, init.AC[i], alg.solver, abs2(z), 1) @@ -184,7 +188,7 @@ function propagator( end end - a = dot(AC_projection(1, init, A, mixedenvs), init.AC[1]) + a = dot(AC_projection(1, init, A, mixedenvs; alg.backend, allocator), init.AC[1]) cb = leftenv(envs1, 1, A) * TransferMatrix(init.AL, H[1:length(A.AL)], A.AL) b = zero(a) for i in 1:length(cb) diff --git a/src/algorithms/statmech/gradient_grassmann.jl b/src/algorithms/statmech/gradient_grassmann.jl index 20bef7b38..8084cbda2 100644 --- a/src/algorithms/statmech/gradient_grassmann.jl +++ b/src/algorithms/statmech/gradient_grassmann.jl @@ -4,7 +4,9 @@ function leading_boundary( alg::GradientGrassmann, envs::MultilineEnvironments = environments(state, operator, state) ) - fg(x) = GrassmannMPS.fg(x, operator, envs) + # read the scheduler here rather than in `fg`, so that the allocator it selects is inferable + scheduler = Defaults.scheduler[] + fg(x) = GrassmannMPS.fg(x, operator, envs; alg.backend, scheduler) x, _, _, _, normgradhistory = optimize( fg, state, alg.method; diff --git a/src/algorithms/statmech/idmrg.jl b/src/algorithms/statmech/idmrg.jl index 39075fd05..8e283dd6e 100644 --- a/src/algorithms/statmech/idmrg.jl +++ b/src/algorithms/statmech/idmrg.jl @@ -1,6 +1,7 @@ function leading_boundary( ψ::MultilineMPS, operator, alg::IDMRG, envs = environments(ψ, operator, ψ) ) + allocator = default_allocator(ψ, SerialScheduler()) log = IterLog("IDMRG") ϵ::Float64 = 2 * alg.tol local iter @@ -13,7 +14,7 @@ function leading_boundary( # left to right sweep for col in 1:size(ψ, 2) - Hac = AC_hamiltonian(col, ψ, operator, ψ, envs) + Hac = AC_hamiltonian(col, ψ, operator, ψ, envs; alg.backend, allocator) _, ψ.AC[:, col] = fixedpoint(Hac, ψ.AC[:, col], :LM, alg_eigsolve) for row in 1:size(ψ, 1) @@ -27,7 +28,7 @@ function leading_boundary( # right to left sweep for col in size(ψ, 2):-1:1 - Hac = AC_hamiltonian(col, ψ, operator, ψ, envs) + Hac = AC_hamiltonian(col, ψ, operator, ψ, envs; alg.backend, allocator) _, ψ.AC[:, col] = fixedpoint(Hac, ψ.AC[:, col], :LM, alg_eigsolve) for row in 1:size(ψ, 1) @@ -64,6 +65,7 @@ end function leading_boundary( ψ::MultilineMPS, operator, alg::IDMRG2, envs = environments(ψ, operator, ψ) ) + allocator = default_allocator(ψ, SerialScheduler()) size(ψ, 2) < 2 && throw(ArgumentError("unit cell should be >= 2")) ϵ::Float64 = 2 * alg.tol log = IterLog("IDMRG2") @@ -78,7 +80,7 @@ function leading_boundary( # sweep from left to right for site in 1:(size(ψ, 2) - 1) ac2 = AC2(ψ, site; kind = :ACAR) - h = AC2_hamiltonian(site, ψ, operator, ψ, envs) + h = AC2_hamiltonian(site, ψ, operator, ψ, envs; alg.backend, allocator) _, ac2′ = fixedpoint(h, ac2, :LM, alg_eigsolve) for row in 1:size(ψ, 1) @@ -102,7 +104,7 @@ function leading_boundary( ψ.AL[:, end] .= ψ.AC[:, end] ./ ψ.C[:, end] ψ.AC[:, 1] .= _mul_tail.(ψ.AL[:, 1], ψ.C[:, 1]) ac2 = AC2(ψ, site; kind = :ALAC) - h = AC2_hamiltonian(site, ψ, operator, ψ, envs) + h = AC2_hamiltonian(site, ψ, operator, ψ, envs; alg.backend, allocator) _, ac2′ = fixedpoint(h, ac2, :LM, alg_eigsolve) for row in 1:size(ψ, 1) @@ -127,7 +129,7 @@ function leading_boundary( # sweep from right to left for site in reverse(1:(size(ψ, 2) - 1)) ac2 = AC2(ψ, site; kind = :ALAC) - h = AC2_hamiltonian(site, ψ, operator, ψ, envs) + h = AC2_hamiltonian(site, ψ, operator, ψ, envs; alg.backend, allocator) _, ac2′ = fixedpoint(h, ac2, :LM, alg_eigsolve) for row in 1:size(ψ, 1) @@ -149,7 +151,7 @@ function leading_boundary( ψ.AC[:, end] .= _mul_front.(ψ.C[:, end - 1], ψ.AR[:, end]) ψ.AR[:, 1] .= _transpose_front.(ψ.C[:, end] .\ _transpose_tail.(ψ.AC[:, 1])) ac2 = AC2(ψ, 0; kind = :ACAR) - h = AC2_hamiltonian(0, ψ, operator, ψ, envs) + h = AC2_hamiltonian(0, ψ, operator, ψ, envs; alg.backend, allocator) _, ac2′ = fixedpoint(h, ac2, :LM, alg_eigsolve) for row in 1:size(ψ, 1) diff --git a/src/algorithms/statmech/vomps.jl b/src/algorithms/statmech/vomps.jl index faa2406cb..9427f5488 100644 --- a/src/algorithms/statmech/vomps.jl +++ b/src/algorithms/statmech/vomps.jl @@ -17,7 +17,7 @@ Used as the `algorithm` argument of [`leading_boundary`](@ref) and [`approximate * [Vanhecke et al. SciPost Phys. Core 4 (2021)](@cite vanhecke2021) """ -@kwdef struct VOMPS{F} <: Algorithm +@kwdef struct VOMPS{F, B} <: Algorithm "tolerance for convergence criterium" tol::Float64 = Defaults.tol @@ -35,6 +35,9 @@ Used as the `algorithm` argument of [`leading_boundary`](@ref) and [`approximate "callback function applied after each iteration, of signature `finalize(iter, ψ, H, envs) -> ψ, envs`" finalize::F = Defaults._finalize + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end # Internal state of the VOMPS algorithm @@ -57,7 +60,7 @@ function dominant_eigsolve( @assert which === :LM "VOMPS only supports the LM eigenvalue problem" log = IterLog("VOMPS") iter = 0 - ϵ = calc_galerkin(mps, operator, mps, envs) + ϵ = calc_galerkin(mps, operator, mps, envs; alg.backend) alg_environments = adapt_solver(alg.alg_environments; iter, g_global = ϵ) recalculate!(envs, mps, operator, mps, alg_environments) @@ -93,7 +96,7 @@ function Base.iterate(it::IterativeSolver{<:VOMPS}, state) mps, envs = it.finalize(state.iter, mps, state.operator, envs)::typeof((mps, envs)) # error criterion - ϵ = calc_galerkin(mps, state.operator, mps, envs) + ϵ = calc_galerkin(mps, state.operator, mps, envs; it.backend) # update state it.state = VOMPSState(mps, state.operator, envs, state.iter + 1, ϵ) @@ -110,10 +113,11 @@ function localupdate_step!( ACs = similar(mps.AC) dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs + allocator = default_allocator(mps, scheduler) tforeach(eachsite(mps); scheduler) do site dst_ACs[site] = _localupdate_vomps_step!( site, mps, state.operator, state.envs; - alg_orth, parallel = false + alg_orth, it.backend, allocator ) return nothing end @@ -122,19 +126,11 @@ function localupdate_step!( end function _localupdate_vomps_step!( - site, mps, operator, envs; parallel::Bool = false, alg_orth = Defaults.alg_orth() + site, mps, operator, envs; alg_orth = Defaults.alg_orth(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) - if !parallel - AC = AC_projection(site, mps, operator, mps, envs) - C = C_projection(site, mps, operator, mps, envs) - return regauge!(AC, C; alg = alg_orth) - end - - local AC, C - @sync begin - @spawn AC = AC_projection(site, mps, operator, mps, envs) - @spawn C = C_projection(site, mps, operator, mps, envs) - end + AC = AC_projection(site, mps, operator, mps, envs; backend, allocator) + C = C_projection(site, mps, operator, mps, envs; backend, allocator) return regauge!(AC, C; alg = alg_orth) end diff --git a/src/algorithms/timestep/bug.jl b/src/algorithms/timestep/bug.jl index 428b6b50e..1ea9edd4d 100644 --- a/src/algorithms/timestep/bug.jl +++ b/src/algorithms/timestep/bug.jl @@ -36,7 +36,7 @@ To restore a maximal dimension of `D`, apply [`changebonds`](@ref) with an [`Svd * [Ceruti et al. BIT Numer. Math. 62 (2022)](@cite ceruti2022) """ -struct BUG{A, O, G, F} <: Algorithm +struct BUG{A, O, G, F, B} <: Algorithm "algorithm used in the exponential solvers" integrator::A @@ -48,14 +48,17 @@ struct BUG{A, O, G, F} <: Algorithm "callback function applied after each iteration, of signature `finalize(t, ψ, H, envs) -> ψ, envs`" finalize::F + + "backend for tensor contractions and index manipulations" + backend::B end function BUG(; integrator = Defaults.alg_expsolve(), alg_orth = Defaults.alg_orth(), trunc = notrunc(), alg_svd = Defaults.alg_svd(), - finalize = Defaults._finalize + finalize = Defaults._finalize, backend = Defaults.backend() ) alg_gauge = _build_inner_gauge(trunc, alg_svd, alg_orth) - return BUG(integrator, alg_orth, alg_gauge, finalize) + return BUG(integrator, alg_orth, alg_gauge, finalize, backend) end # `ψ.AC[site]` first, the neighbour second: the lazy `CView` walk keys off what is already cached, and @@ -120,19 +123,19 @@ function _augment_basis!(site::Int, dir::Val{:left}, ψ, AR_old_tail, AC, C₀, ) end -function _evolve_center(site, ψ, H, alg::BUG, envs, t, h; imaginary_evolution) - Heff = AC_hamiltonian(site, ψ, H, ψ, envs) +function _evolve_center(site, ψ, H, alg::BUG, envs, t, h, allocator; imaginary_evolution) + Heff = AC_hamiltonian(site, ψ, H, ψ, envs; alg.backend, allocator) return integrate(Heff, ψ.AC[site], t, h, alg.integrator; imaginary_evolution) end function local_update!( - site, direction::Val, ψ, H, alg::BUG, envs, t, h; + site, direction::Val, ψ, H, alg::BUG, envs, t, h, allocator; imaginary_evolution, normalize, timeroutput ) # at the far end of the sweep there is no bond ahead to cut: evolve and finalize if site == _sweep_end(ψ, direction) AC = @timeit timeroutput "AC_integrate" _evolve_center( - site, ψ, H, alg, envs, t, h; imaginary_evolution + site, ψ, H, alg, envs, t, h, allocator; imaginary_evolution ) normalize && normalize!(AC) ψ.AC[site] = AC @@ -146,7 +149,7 @@ function local_update!( # 2. evolve the connecting tensor AC = @timeit timeroutput "AC_integrate" _evolve_center( - site, ψ, H, alg, envs, t, h; imaginary_evolution + site, ψ, H, alg, envs, t, h, allocator; imaginary_evolution ) # 3. augment the basis (old first, no truncation here) and install it, together with the @@ -165,10 +168,13 @@ function timestep!( L = length(ψ) h = dt / 2 + # the sweep is serial, so a single allocator serves all local updates + allocator = default_allocator(ψ, SerialScheduler()) + # left→right half-sweep (root = last site): `t → t + dt / 2` @timeit timeroutput "half-sweep" for site in 1:L ψ = local_update!( - site, Val(:right), ψ, H, alg, envs, t, h; + site, Val(:right), ψ, H, alg, envs, t, h, allocator; imaginary_evolution, normalize, timeroutput ) end @@ -176,7 +182,7 @@ function timestep!( # right→left half-sweep (root = first site): `t + dt / 2 → t + dt` @timeit timeroutput "half-sweep" for site in L:-1:1 ψ = local_update!( - site, Val(:left), ψ, H, alg, envs, t + h, h; + site, Val(:left), ψ, H, alg, envs, t + h, h, allocator; imaginary_evolution, normalize, timeroutput ) end diff --git a/src/algorithms/timestep/tdvp.jl b/src/algorithms/timestep/tdvp.jl index 9a4fbf533..a3da18c52 100644 --- a/src/algorithms/timestep/tdvp.jl +++ b/src/algorithms/timestep/tdvp.jl @@ -29,7 +29,7 @@ Used as the `algorithm` argument of [`timestep`](@ref), [`timestep!`](@ref) and * [Haegeman et al. Phys. Rev. Lett. 107 (2011)](@cite haegeman2011) """ -struct TDVP{A, E, G, F} <: Algorithm +struct TDVP{A, E, G, F, B} <: Algorithm "algorithm used in the exponential solvers" integrator::A @@ -47,12 +47,16 @@ struct TDVP{A, E, G, F} <: Algorithm "callback function applied after each iteration, of signature `finalize(t, ψ, H, envs) -> ψ, envs`" finalize::F + + "backend for tensor contractions and index manipulations" + backend::B end function TDVP(; integrator = Defaults.alg_expsolve(), tolgauge = Defaults.tolgauge, gaugemaxiter = Defaults.maxiter, finalize = Defaults._finalize, alg_expand = nothing, trunc = notrunc(), - alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth() + alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth(), + backend = Defaults.backend() ) # a no-truncation `trunc` selects a (bond-preserving) QR gauge, anything else a truncated SVD alg_gauge = trunc isa MatrixAlgebraKit.NoTruncation ? alg_orth : @@ -60,7 +64,9 @@ function TDVP(; if !isnothing(alg_expand) && !_truncates(alg_gauge) @warn "TDVP with `alg_expand` but no truncation (`trunc = notrunc()`): the bond dimension will grow unboundedly each sweep." end - return TDVP(integrator, tolgauge, gaugemaxiter, alg_expand, alg_gauge, finalize) + return TDVP( + integrator, tolgauge, gaugemaxiter, alg_expand, alg_gauge, finalize, backend + ) end function timestep( @@ -76,39 +82,40 @@ function timestep( return timestep(complex(ψ), H, t, dt, alg, envs; leftorthflag, imaginary_evolution, normalize) end + # the scheduler is read here rather than below, so that the allocator it selects is inferable + return _timestep_infinite( + ψ, H, t, dt, alg, envs, Defaults.scheduler[]; leftorthflag, imaginary_evolution + ) +end + +function _timestep_infinite( + ψ::InfiniteMPS, H, t::Number, dt::Number, alg::TDVP, envs, scheduler::Scheduler; + leftorthflag, imaginary_evolution + ) temp_ACs = similar(ψ.AC) temp_Cs = similar(ψ.C) - scheduler = Defaults.scheduler[] + # both sweeps together are a single unit of concurrent work, and share one allocator + allocator = default_allocator(ψ, scheduler) + ac_sweep!() = tforeach(1:length(ψ); scheduler) do loc + Hac = AC_hamiltonian(loc, ψ, H, ψ, envs; alg.backend, allocator) + temp_ACs[loc] = integrate(Hac, ψ.AC[loc], t, dt, alg.integrator; imaginary_evolution) + return nothing + end + c_sweep!() = tforeach(1:length(ψ); scheduler) do loc + Hc = C_hamiltonian(loc, ψ, H, ψ, envs; alg.backend, allocator) + temp_Cs[loc] = integrate(Hc, ψ.C[loc], t, dt, alg.integrator; imaginary_evolution) + return nothing + end + if scheduler isa SerialScheduler - temp_ACs = tmap!(temp_ACs, 1:length(ψ); scheduler) do loc - Hac = AC_hamiltonian(loc, ψ, H, ψ, envs) - return integrate(Hac, ψ.AC[loc], t, dt, alg.integrator; imaginary_evolution) - end - temp_Cs = tmap!(temp_Cs, 1:length(ψ); scheduler) do loc - Hc = C_hamiltonian(loc, ψ, H, ψ, envs) - return integrate(Hc, ψ.C[loc], t, dt, alg.integrator; imaginary_evolution) - end + ac_sweep!() + c_sweep!() else + # the AC and C sweeps are independent, so run them concurrently with each other too @sync begin - Threads.@spawn begin - temp_ACs = tmap!(temp_ACs, 1:length(ψ); scheduler) do loc - Hac = AC_hamiltonian(loc, ψ, H, ψ, envs) - return integrate( - Hac, ψ.AC[loc], t, dt, alg.integrator; - imaginary_evolution - ) - end - end - Threads.@spawn begin - temp_Cs = tmap!(temp_Cs, 1:length(ψ); scheduler) do loc - Hc = C_hamiltonian(loc, ψ, H, ψ, envs) - return integrate( - Hc, ψ.C[loc], t, dt, alg.integrator; - imaginary_evolution - ) - end - end + Threads.@spawn ac_sweep!() + Threads.@spawn c_sweep!() end end @@ -130,15 +137,25 @@ function timestep!( envs::AbstractMPSEnvironments = environments(ψ, H, ψ); imaginary_evolution::Bool = false, normalize::Bool = false ) + # the sweep is serial, so a single allocator serves all local updates + allocator = default_allocator(ψ, SerialScheduler()) + return _timestep_finite!( + ψ, H, t, dt, alg, envs, allocator; imaginary_evolution, normalize + ) +end +function _timestep_finite!( + ψ::AbstractFiniteMPS, H, t::Number, dt::Number, alg::TDVP, envs, allocator; + imaginary_evolution::Bool, normalize::Bool + ) # sweep left to right for i in 1:(length(ψ) - 1) # 1. optionally expand the bond ahead of the local update (CBE) isnothing(alg.alg_expand) || - changebond!(i, Val(:right), ψ, H, alg.alg_expand, envs; normalize) + changebond!(i, Val(:right), ψ, H, alg.alg_expand, envs; normalize, allocator) # 2. evolve the (possibly expanded) center tensor forward - Hac = AC_hamiltonian(i, ψ, H, ψ, envs) + Hac = AC_hamiltonian(i, ψ, H, ψ, envs; alg.backend, allocator) AC = integrate(Hac, ψ.AC[i], t, dt / 2, alg.integrator; imaginary_evolution) # 3. gauge: split AC -> AL[i], C[i] (QR center-move, or truncated SVD cutting the @@ -147,7 +164,7 @@ function timestep!( left_gauge!(ψ, i, AC, alg.alg_gauge; normalize) # 4. evolve the bond tensor backward - Hc = C_hamiltonian(i, ψ, H, ψ, envs) + Hc = C_hamiltonian(i, ψ, H, ψ, envs; alg.backend, allocator) ψ.C[i] = integrate( Hc, ψ.C[i], t + dt / 2, -dt / 2, alg.integrator; imaginary_evolution @@ -155,17 +172,17 @@ function timestep!( end # edge case - Hac = AC_hamiltonian(length(ψ), ψ, H, ψ, envs) + Hac = AC_hamiltonian(length(ψ), ψ, H, ψ, envs; alg.backend, allocator) ψ.AC[end] = integrate(Hac, ψ.AC[end], t, dt / 2, alg.integrator; imaginary_evolution) # sweep right to left for i in length(ψ):-1:2 # 1. optionally expand the bond ahead of the local update (CBE) isnothing(alg.alg_expand) || - changebond!(i, Val(:left), ψ, H, alg.alg_expand, envs; normalize) + changebond!(i, Val(:left), ψ, H, alg.alg_expand, envs; normalize, allocator) # 2. evolve the (possibly expanded) center tensor forward - Hac = AC_hamiltonian(i, ψ, H, ψ, envs) + Hac = AC_hamiltonian(i, ψ, H, ψ, envs; alg.backend, allocator) AC = integrate( Hac, ψ.AC[i], t + dt / 2, dt / 2, alg.integrator; imaginary_evolution @@ -176,7 +193,7 @@ function timestep!( right_gauge!(ψ, i, AC, alg.alg_gauge; normalize) # 4. evolve the bond tensor backward - Hc = C_hamiltonian(i - 1, ψ, H, ψ, envs) + Hc = C_hamiltonian(i - 1, ψ, H, ψ, envs; alg.backend, allocator) ψ.C[i - 1] = integrate( Hc, ψ.C[i - 1], t + dt, -dt / 2, alg.integrator; imaginary_evolution @@ -184,7 +201,7 @@ function timestep!( end # edge case - Hac = AC_hamiltonian(1, ψ, H, ψ, envs) + Hac = AC_hamiltonian(1, ψ, H, ψ, envs; alg.backend, allocator) ψ.AC[1] = integrate( Hac, ψ.AC[1], t + dt / 2, dt / 2, alg.integrator; imaginary_evolution @@ -210,7 +227,7 @@ Used as the `algorithm` argument of [`timestep`](@ref), [`timestep!`](@ref) and * [Haegeman et al. Phys. Rev. Lett. 107 (2011)](@cite haegeman2011) """ -@kwdef struct TDVP2{A, S, F} <: Algorithm +@kwdef struct TDVP2{A, S, F, B} <: Algorithm "algorithm used in the exponential solvers" integrator::A = Defaults.alg_expsolve() @@ -228,6 +245,9 @@ Used as the `algorithm` argument of [`timestep`](@ref), [`timestep!`](@ref) and "callback function applied after each iteration, of signature `finalize(t, ψ, H, envs) -> ψ, envs`" finalize::F = Defaults._finalize + + "backend for tensor contractions and index manipulations" + backend::B = Defaults.backend() end function timestep!( @@ -235,11 +255,21 @@ function timestep!( envs::AbstractMPSEnvironments = environments(ψ, H, ψ); imaginary_evolution::Bool = false, normalize::Bool = false ) + # the sweep is serial, so a single allocator serves all local updates + allocator = default_allocator(ψ, SerialScheduler()) + return _timestep2_finite!( + ψ, H, t, dt, alg, envs, allocator; imaginary_evolution, normalize + ) +end +function _timestep2_finite!( + ψ::AbstractFiniteMPS, H, t::Number, dt::Number, alg::TDVP2, envs, allocator; + imaginary_evolution::Bool, normalize::Bool + ) # sweep left to right for i in 1:(length(ψ) - 1) ac2 = _transpose_front(ψ.AC[i]) * _transpose_tail(ψ.AR[i + 1]) - Hac2 = AC2_hamiltonian(i, ψ, H, ψ, envs) + Hac2 = AC2_hamiltonian(i, ψ, H, ψ, envs; alg.backend, allocator) ac2′ = integrate(Hac2, ac2, t, dt / 2, alg.integrator; imaginary_evolution) nal, nc, nar = svd_trunc!(ac2′; trunc = alg.trunc, alg = alg.alg_svd) @@ -248,7 +278,7 @@ function timestep!( ψ.AC[i + 1] = (complex(nc), _transpose_front(nar)) if i != (length(ψ) - 1) - Hac = AC_hamiltonian(i + 1, ψ, H, ψ, envs) + Hac = AC_hamiltonian(i + 1, ψ, H, ψ, envs; alg.backend, allocator) ψ.AC[i + 1] = integrate( Hac, ψ.AC[i + 1], t + dt / 2, -dt / 2, alg.integrator; imaginary_evolution @@ -259,7 +289,7 @@ function timestep!( # sweep right to left for i in length(ψ):-1:2 ac2 = _transpose_front(ψ.AL[i - 1]) * _transpose_tail(ψ.AC[i]) - Hac2 = AC2_hamiltonian(i - 1, ψ, H, ψ, envs) + Hac2 = AC2_hamiltonian(i - 1, ψ, H, ψ, envs; alg.backend, allocator) ac2′ = integrate(Hac2, ac2, t + dt / 2, dt / 2, alg.integrator; imaginary_evolution) nal, nc, nar = svd_trunc!(ac2′; trunc = alg.trunc, alg = alg.alg_svd) @@ -268,7 +298,7 @@ function timestep!( ψ.AC[i] = (complex(nc), _transpose_front(nar)) if i != 2 - Hac = AC_hamiltonian(i - 1, ψ, H, ψ, envs) + Hac = AC_hamiltonian(i - 1, ψ, H, ψ, envs; alg.backend, allocator) ψ.AC[i - 1] = integrate( Hac, ψ.AC[i - 1], t + dt, -dt / 2, alg.integrator; imaginary_evolution diff --git a/src/algorithms/timestep/time_evolve.jl b/src/algorithms/timestep/time_evolve.jl index 85193791d..b4a1f8043 100644 --- a/src/algorithms/timestep/time_evolve.jl +++ b/src/algorithms/timestep/time_evolve.jl @@ -37,7 +37,7 @@ for (timestep, time_evolve) in zip((:timestep, :timestep!), (:time_evolve, :time ) log = IterLog("TDVP") LoggingExtras.withlevel(; verbosity) do - @infov 2 loginit!(log, 0, t) + @infov 2 loginit!(log, 0, first(t_span)) for iter in 1:(length(t_span) - 1) t = t_span[iter] dt = t_span[iter + 1] - t diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl index de11eabcd..b995329f1 100644 --- a/src/algorithms/toolbox.jl +++ b/src/algorithms/toolbox.jl @@ -43,8 +43,8 @@ function infinite_temperature_density_matrix(H::MPOHamiltonian) end """ - calc_galerkin(below, operator, above, envs) - calc_galerkin(pos, below, operator, above, envs) + calc_galerkin(below, operator, above, envs; kwargs...) + calc_galerkin(pos, below, operator, above, envs; kwargs...) Calculate the Galerkin error, which is the error between the solution of the original problem, and the solution of the problem projected on the tangent space. Concretely, this is the overlap of the current state with the single-site derivative, projected onto the nullspace of the current state: @@ -52,18 +52,29 @@ Concretely, this is the overlap of the current state with the single-site deriva ```math \\epsilon = \\left|VL ⋅ \\left(VL^{\\dagger} ⋅ \\frac{\\partial \\text{above}}{\\partial AC_{\\text{pos}}}\\right)\\right| ``` + +# Keyword Arguments + +- `backend = DefaultBackend()`: backend for the tensor contractions of the derivative. +- `allocator = DefaultAllocator()`: allocator serving their scratch space. A sweep that already + holds one should pass it, rather than leaving this contraction to the garbage collector. """ -function calc_galerkin(pos::Int, below, operator, above, envs) - AC´ = AC_projection(pos, below, operator, above, envs) +function calc_galerkin( + pos::Int, below, operator, above, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + AC´ = AC_projection(pos, below, operator, above, envs; backend, allocator) normalize!(AC´) return norm(project_complement!(AC´, below.AL[pos])) end -function calc_galerkin(pos::CartesianIndex{2}, below, operator, above, envs) +function calc_galerkin(pos::CartesianIndex{2}, below, operator, above, envs; kwargs...) row, col = Tuple(pos) - return calc_galerkin(col, below[row + 1], operator[row], above[row], envs[row]) + return calc_galerkin(col, below[row + 1], operator[row], above[row], envs[row]; kwargs...) end -function calc_galerkin(below, operator, above, envs) - return maximum(pos -> calc_galerkin(pos, below, operator, above, envs), eachindex(below)) +function calc_galerkin(below, operator, above, envs; kwargs...) + return maximum(eachindex(below)) do pos + return calc_galerkin(pos, below, operator, above, envs; kwargs...) + end end """ diff --git a/src/utility/allocator.jl b/src/utility/allocator.jl new file mode 100644 index 000000000..3dc4d64f2 --- /dev/null +++ b/src/utility/allocator.jl @@ -0,0 +1,41 @@ +# Scratch space for local updates +# ------------------------------- +# The local updates of an MPS algorithm allocate a fair number of intermediates, all of which die +# before the next update. Serving them from an allocator rather than from Julia's memory manager +# keeps them out of the garbage collector's way; which allocator is appropriate depends on where the +# tensors live and on whether more than one task will be sharing it. + +""" + default_allocator(x, scheduler) -> allocator + +The allocator that serves the scratch space of local updates on `x`, for work scheduled with +`scheduler`. + +`x` is anything with a `storagetype`, typically the state being operated on. Host memory gets a +`TensorOperations.BufferAllocator`, which serves intermediates from a reusable buffer, when a single +task owns the allocator, and a `TensorOperations.ManualAllocator`, which `malloc`s and `free`s them +one by one, when the allocator is shared between tasks - a buffer is not thread-safe, whereas a +manual allocator holds no state at all. Any other storage type falls back on +`TensorOperations.DefaultAllocator`, which allocates through the storage type itself and is +therefore correct on any device, at the cost of leaving intermediates to the garbage collector. + +Extend this function to serve a storage type that MPSKit does not know about. Dedicated scratch space +can be turned off altogether with `MPSKit.Defaults.set_buffering!`. + +!!! warning + An allocator obtained for a `SerialScheduler` must not be shared between tasks. Sites that spawn + should pass the scheduler they spawn with, so that the allocator matches the concurrency. +""" +default_allocator(x, scheduler::Scheduler) = default_allocator(storagetype(x), scheduler) + +# `Memory` only exists from Julia 1.11, and `BufferAllocator`'s own storage follows suit, so the +# host-storage test is taken from TensorOperations rather than hardcoded +const HostStorage = @static isdefined(Core, :Memory) ? Union{Array, Memory} : Array + +# The scheduler is what knows whether the allocator will be shared, so it - rather than the author of +# any individual sweep - decides between a buffer and a manual allocator. +if Defaults.buffering + default_allocator(::Type{<:HostStorage}, ::SerialScheduler) = BufferAllocator() + default_allocator(::Type{<:HostStorage}, ::Scheduler) = ManualAllocator() +end +default_allocator(::Type, ::Scheduler) = DefaultAllocator() diff --git a/src/utility/defaults.jl b/src/utility/defaults.jl index dbc221420..c6c100c29 100644 --- a/src/utility/defaults.jl +++ b/src/utility/defaults.jl @@ -8,7 +8,9 @@ module Defaults import KrylovKit: GMRES, Arnoldi, Lanczos using OhMyThreads using ..MPSKit: DynamicTol, AdaptiveKrylov +using Preferences: @load_preference, @set_preferences! using TensorKit: TensorKit +using TensorOperations: DefaultBackend using MatrixAlgebraKit: DefaultAlgorithm, Householder const VERBOSE_NONE = 0 @@ -41,6 +43,17 @@ const eigsolver = Arnoldi(; tol, maxiter, eager = true) alg_svd() = DefaultAlgorithm() alg_orth() = Householder(; positive = true) +""" + backend() + +The default backend for tensor contractions and index manipulations. + +`TensorOperations.DefaultBackend` is a placeholder rather than a CPU backend: the actual +implementation is selected from the types of the tensors involved, so this is already the right +choice for GPU-backed states. +""" +backend() = DefaultBackend() + function alg_gauge(; tol = tolgauge, maxiter = maxiter, verbosity = VERBOSE_WARN, alg_orth = alg_orth(), @@ -82,7 +95,7 @@ end """ const scheduler -A scoped value that controls the current settings for multi-threading, typically used to parallelize over unitcells. +The current settings for multi-threading, typically used to parallelize over unitcells. This value is best controlled using [`set_scheduler!`](@ref). """ const scheduler = Ref{Scheduler}() @@ -105,4 +118,36 @@ function set_scheduler!(sc = OhMyThreads.Implementation.NotGiven(); kwargs...) return scheduler[] end +""" + const buffering + +Whether local updates serve their scratch space from a dedicated allocator. + +This is a compile-time preference, best controlled using [`set_buffering!`](@ref). +""" +const buffering = @load_preference("buffering", true)::Bool + +""" + set_buffering!(b::Bool) + +Enable or disable dedicated scratch space for local updates. + +When enabled - the default - the tensor contractions of a local update serve their intermediates +from an allocator that bypasses Julia's memory manager, which cuts both the allocation count and the +garbage-collection time substantially. Disabling it trades that back for lower memory use, which is +worthwhile when memory rather than time is the binding constraint. + +This only affects storage types for which MPSKit has such an allocator, i.e. host memory; see +[`MPSKit.default_allocator`](@ref). + +!!! note + This setting is stored in a `LocalPreferences.toml` file next to the active `Project.toml` and is + read when MPSKit is compiled, so Julia has to be restarted for a change to take effect. +""" +function set_buffering!(b::Bool) + @set_preferences!("buffering" => b) + @info "Buffering set to $b; restart Julia for this to take effect." + return b +end + end diff --git a/test/algorithms/allocator.jl b/test/algorithms/allocator.jl new file mode 100644 index 000000000..02ba76f12 --- /dev/null +++ b/test/algorithms/allocator.jl @@ -0,0 +1,156 @@ +println(" +------------------------------------ +| Allocators in algorithms | +------------------------------------ +") + +using .TestSetup +using Test, TestExtras +using MPSKit +using TensorKit +using TensorKit: ℙ +using MPSKit: SerialScheduler, DynamicScheduler +using Random + +verbosity = 0 + +# The scheduler is process-global: always restore it. +function with_scheduler(f, scheduler) + old_sched = MPSKit.Defaults.scheduler[] + try + MPSKit.Defaults.scheduler[] = scheduler + return f() + finally + MPSKit.Defaults.scheduler[] = old_sched + end +end + +# The scheduler decides which allocator the local updates get - a buffer when a single task owns it, +# a shared manual allocator otherwise - so it must not change any number. +const SCHEDULERS = ("serial" => SerialScheduler(), "dynamic" => DynamicScheduler()) + +@testset "finite algorithms" verbose = true begin + L = 8 + H = force_planar(transverse_field_ising(; g = 4.0, L)) + Random.seed!(0x1234) + ψ₀ = FiniteMPS(randn, ComplexF64, L, ℙ^2, ℙ^8) + + @testset "DMRG" begin + alg = DMRG(; verbosity, maxiter = 10) + ψ, _, _ = find_groundstate(ψ₀, H, alg) + E = expectation_value(ψ, H) + + # the algorithm object holds no scratch space, so a second solve reproduces the first + ψ2, _, _ = find_groundstate(ψ₀, H, alg) + @test expectation_value(ψ2, H) ≈ E atol = 1.0e-8 + end + + @testset "DMRG2" begin + reference = expectation_value(first(find_groundstate(ψ₀, H, DMRG(; verbosity, maxiter = 10))), H) + alg = DMRG2(; verbosity, maxiter = 10, trunc = truncrank(8)) + ψ, _, _ = find_groundstate(ψ₀, H, alg) + @test expectation_value(ψ, H) ≈ reference atol = 1.0e-6 + end + + @testset "TDVP" begin + ψ, _ = timestep(complex(ψ₀), H, 0.0, 0.01, TDVP()) + @test norm(ψ) ≈ norm(ψ₀) atol = 1.0e-6 + end + + @testset "BUG" begin + # the allocator only serves the effective-Hamiltonian applications, so a buffer-served + # step has to land on the same state as one that is not + ψ, _ = timestep(complex(ψ₀), H, 0.0, 0.01, BUG()) + @test norm(ψ) ≈ norm(ψ₀) atol = 1.0e-6 + + ψ_tdvp, _ = timestep(complex(ψ₀), H, 0.0, 0.01, TDVP()) + @test abs(dot(ψ, ψ_tdvp)) / (norm(ψ) * norm(ψ_tdvp)) ≈ 1.0 atol = 1.0e-4 + end + + @testset "nested algorithms" begin + # `alg_expand` is driven through `changebond!`, which is handed the sweep's allocator + # rather than selecting one of its own once per site + alg = DMRG(; + verbosity, maxiter = 3, alg_expand = SketchedExpand(; trunc = truncrank(2)), + trunc = truncrank(8), + ) + ψ, _, _ = find_groundstate(ψ₀, H, alg) + @test expectation_value(ψ, H) isa Number + end + + @testset "time_evolve reuses nothing across steps" begin + nsteps = Ref(0) + finalize = function (t, ψ, H, envs) + nsteps[] += 1 + return ψ, envs + end + ψ, _ = time_evolve(complex(ψ₀), H, 0.0:0.01:0.03, TDVP(; finalize)) + @test nsteps[] == 3 + @test norm(ψ) ≈ norm(ψ₀) atol = 1.0e-6 + end +end + +@testset "infinite algorithms" verbose = true begin + H = force_planar(repeat(transverse_field_ising(; g = 4.0), 2)) + Random.seed!(0x1234) + ψ₀ = InfiniteMPS([ℙ^2, ℙ^2], [ℙ^8, ℙ^8]) + + reference = nothing + for (schedname, scheduler) in SCHEDULERS + @testset "VUMPS ($schedname)" begin + alg = VUMPS(; verbosity, maxiter = 50) + ψ, _, _ = with_scheduler(scheduler) do + find_groundstate(ψ₀, H, alg) + end + E = sum(expectation_value(ψ, H)) + isnothing(reference) ? (reference = E) : (@test E ≈ reference atol = 1.0e-6) + end + + @testset "GradientGrassmann ($schedname)" begin + alg = GradientGrassmann(; verbosity, maxiter = 50) + ψ, _, _ = with_scheduler(scheduler) do + find_groundstate(ψ₀, H, alg) + end + @test sum(expectation_value(ψ, H)) ≈ reference atol = 1.0e-4 + end + + @testset "TDVP ($schedname)" begin + ψ, _ = with_scheduler(scheduler) do + timestep(ψ₀, H, 0.0, 0.01, TDVP()) + end + @test sum(expectation_value(ψ, H)) isa Number + end + + @testset "VOMPS ($schedname)" begin + # leading_boundary and approximate take separate code paths, and the latter is the one + # that runs the AC and C projections of a site concurrently with each other + β = 0.5 + O = classical_ising(; β) + ψ, envs = with_scheduler(scheduler) do + leading_boundary(InfiniteMPS(ℂ^2, ℂ^10), O, VOMPS(; verbosity, tol = 1.0e-8)) + end + @test -log(expectation_value(ψ, O, envs)) / β ≈ -2.0515856253898357 atol = 1.0e-8 + + dt = 1.0e-3 + W = make_time_mpo(H, dt, WII()) + ϕ = InfiniteMPS([ℙ^2, ℙ^2], [ℙ^10, ℙ^10]) + ϕ′, _ = with_scheduler(scheduler) do + approximate(InfiniteMPS([ℙ^2, ℙ^2], [ℙ^12, ℙ^12]), (W, ϕ), VOMPS(; verbosity)) + end + ϕ_ref, _ = timestep(ϕ, H, 0.0, dt, TDVP()) + @test abs(dot(ϕ′, ϕ_ref)) ≈ 1.0 atol = dt + end + end + + @testset "IDMRG" begin + alg = IDMRG(; verbosity, maxiter = 100) + ψ, _, _ = find_groundstate(ψ₀, H, alg) + @test sum(expectation_value(ψ, H)) ≈ reference atol = 1.0e-4 + end + + @testset "changebonds" begin + alg = OptimalExpand(; trunc = truncrank(4)) + ψ, _ = changebonds(ψ₀, H, alg) + @test dim(left_virtualspace(ψ, 1)) > dim(left_virtualspace(ψ₀, 1)) + end +end diff --git a/test/gpu/amd/algorithms.jl b/test/gpu/amd/algorithms.jl new file mode 100644 index 000000000..313e80226 --- /dev/null +++ b/test/gpu/amd/algorithms.jl @@ -0,0 +1,8 @@ +using Test, TestExtras +using AMDGPU + +const ArrType = ROCArray + +@testset "AMDGPU algorithms" verbose = true begin + include(joinpath(@__DIR__, "..", "..", "setup", "gpu_algorithms.jl")) +end diff --git a/test/gpu/cuda/algorithms.jl b/test/gpu/cuda/algorithms.jl new file mode 100644 index 000000000..63bffec47 --- /dev/null +++ b/test/gpu/cuda/algorithms.jl @@ -0,0 +1,8 @@ +using Test, TestExtras +using CUDA, cuTENSOR + +const ArrType = CuArray + +@testset "CUDA algorithms" verbose = true begin + include(joinpath(@__DIR__, "..", "..", "setup", "gpu_algorithms.jl")) +end diff --git a/test/misc/allocator.jl b/test/misc/allocator.jl new file mode 100644 index 000000000..5cd9775a2 --- /dev/null +++ b/test/misc/allocator.jl @@ -0,0 +1,134 @@ +println(" +--------------------------- +| Allocators | +--------------------------- +") + +using .TestSetup +using Test, TestExtras +using MPSKit +using TensorKit +using MPSKit: default_allocator, BufferAllocator, ManualAllocator, DefaultAllocator, HostStorage +using MPSKit: AC_hamiltonian, environments +# `OhMyThreads` is not a test dependency of its own, so take the threading names from MPSKit +using MPSKit: Scheduler, SerialScheduler, DynamicScheduler, StaticScheduler, GreedyScheduler, + tforeach + +# stand-in for a storage type MPSKit has no dedicated allocator for, e.g. a device array +struct FakeDeviceVector{T} <: DenseVector{T} end + +# `Base.infer_return_type` is Julia 1.11+; `Core.Compiler.return_type` covers the LTS as well, +# and is what MPSKit itself already uses (e.g. `grassmann.jl`). +infer_rt(f, types::Tuple) = @static if isdefined(Base, :infer_return_type) + Base.infer_return_type(f, types) +else + Core.Compiler.return_type(f, Tuple{types...}) +end + +const SCHEDULERS = ( + "dynamic" => DynamicScheduler(), + "dynamic nchunks" => DynamicScheduler(; nchunks = 3), + "dynamic chunksize" => DynamicScheduler(; chunksize = 2), + "dynamic unchunked" => DynamicScheduler(; chunking = false), + "dynamic interactive" => DynamicScheduler(; threadpool = :interactive), + "static" => StaticScheduler(), + "greedy" => GreedyScheduler(), +) + +@testset "dispatch on storage type" begin + # host storage gets a dedicated allocator, anything else falls back on the storage type itself + @test default_allocator(Vector{ComplexF64}, SerialScheduler()) isa BufferAllocator + @static if isdefined(Core, :Memory) + @test Memory{Float64} <: HostStorage + @test default_allocator(Memory{Float64}, SerialScheduler()) isa BufferAllocator + end + @test default_allocator(FakeDeviceVector{ComplexF64}, SerialScheduler()) isa DefaultAllocator + @test default_allocator(FakeDeviceVector{ComplexF64}, DynamicScheduler()) isa DefaultAllocator + + # a state is dispatched on through its storage type + ψ = FiniteMPS(randn, ComplexF64, 4, ℂ^2, ℂ^4) + @test default_allocator(ψ, SerialScheduler()) isa BufferAllocator + @test default_allocator(TensorKit.storagetype(ψ), SerialScheduler()) isa BufferAllocator +end + +@testset "a shared allocator is never a buffer ($name)" for (name, scheduler) in SCHEDULERS + # `BufferAllocator` has a mutable offset and is not thread-safe, so it must never be handed to a + # scheduler that spawns. This is the invariant the whole design rests on. + ψ = FiniteMPS(randn, ComplexF64, 4, ℂ^2, ℂ^4) + allocator = default_allocator(ψ, scheduler) + @test allocator isa ManualAllocator + @test !(allocator isa BufferAllocator) +end + +@testset "inference" begin + ψ = FiniteMPS(randn, ComplexF64, 4, ℂ^2, ℂ^4) + S = TensorKit.storagetype(ψ) + + # The allocator reaching the contractions has to be concrete, or every one of them goes dynamic. + # For a statically known scheduler that holds exactly: + @test isconcretetype(infer_rt(default_allocator, (typeof(ψ), SerialScheduler))) + @test infer_rt(default_allocator, (Type{S}, SerialScheduler)) === typeof(BufferAllocator()) + @test infer_rt(default_allocator, (typeof(ψ), typeof(DynamicScheduler()))) === ManualAllocator + @test @constinferred(default_allocator(ψ, SerialScheduler())) isa BufferAllocator + + # `Defaults.scheduler[]` is abstractly typed, so a site that does not pass its scheduler through + # a function boundary gets a small union instead - which is why the algorithms do pass it. + rt = infer_rt(default_allocator, (typeof(ψ), Scheduler)) + @test !isconcretetype(rt) + @test typeof(BufferAllocator()) <: rt && ManualAllocator <: rt +end + +@testset "buffering preference" begin + # A compile-time preference cannot be flipped in-process, so this only checks that the selector + # agrees with whatever MPSKit was compiled with. + ψ = FiniteMPS(randn, ComplexF64, 4, ℂ^2, ℂ^4) + if MPSKit.Defaults.buffering + @test default_allocator(ψ, SerialScheduler()) isa BufferAllocator + @test default_allocator(ψ, DynamicScheduler()) isa ManualAllocator + else + @test default_allocator(ψ, SerialScheduler()) isa DefaultAllocator + @test default_allocator(ψ, DynamicScheduler()) isa DefaultAllocator + end +end + +@testset "the three allocators agree" begin + # Every allocator has to produce the same numbers; the shared one in particular serves its + # intermediates from `PtrArray`s rather than `Vector`s, which is a different code path in + # TensorKit. + L = 6 + H = transverse_field_ising(; g = 1.5, L) + ψ = FiniteMPS(randn, ComplexF64, L, ℂ^2, ℂ^8) + envs = environments(ψ, H, ψ) + + reference = AC_hamiltonian(3, ψ, H, ψ, envs; allocator = DefaultAllocator())(ψ.AC[3]) + for allocator in (BufferAllocator(), ManualAllocator()) + y = AC_hamiltonian(3, ψ, H, ψ, envs; allocator)(ψ.AC[3]) + @test y ≈ reference + end +end + +@testset "a shared allocator survives concurrent use" begin + # The threaded path hands one allocator to every task at once, so it has to hold no state that + # could be corrupted by that. + L = 6 + H = transverse_field_ising(; g = 1.5, L) + ψ = FiniteMPS(randn, ComplexF64, L, ℂ^2, ℂ^8) + envs = environments(ψ, H, ψ) + + sites = 2:(L - 1) + reference = map(sites) do i + return AC_hamiltonian(i, ψ, H, ψ, envs; allocator = DefaultAllocator())(ψ.AC[i]) + end + + scheduler = DynamicScheduler(; chunking = false) + allocator = default_allocator(ψ, scheduler) + out = Vector{Any}(undef, length(sites)) + for _ in 1:20 + tforeach(eachindex(sites); scheduler) do j + H_ac = AC_hamiltonian(sites[j], ψ, H, ψ, envs; allocator) + out[j] = H_ac(ψ.AC[sites[j]]) + return nothing + end + @test all(out[j] ≈ reference[j] for j in eachindex(sites)) + end +end diff --git a/test/setup/gpu_algorithms.jl b/test/setup/gpu_algorithms.jl new file mode 100644 index 000000000..e07e5eb1d --- /dev/null +++ b/test/setup/gpu_algorithms.jl @@ -0,0 +1,105 @@ +# Shared body for the CUDA and AMDGPU algorithm tests. +# +# The thin per-vendor wrappers in `test/gpu/{cuda,amd}/algorithms.jl` set `ArrType` (e.g. `CuArray`), +# then `include` this file. It lives under `setup/` because `runtests.jl` filters that prefix out of +# test discovery - anywhere under `test/gpu/` it would be picked up and run on its own, with +# `ArrType` undefined. +# +# This is the coverage that was missing entirely: `test/gpu/` exercised construction, state and +# operator algebra, and `expectation_value`, but never ran an algorithm, so nothing checked that the +# scratch-space allocator handed to a local update actually matches the device the state lives on. + +using MPSKit +using MPSKit: default_allocator, SerialScheduler, DynamicScheduler +using MPSKit: BufferAllocator, ManualAllocator, DefaultAllocator +using TensorKit +using Adapt +using Random + +# `Base.infer_return_type` is Julia 1.11+; `Core.Compiler.return_type` covers the LTS as well, +# and is what MPSKit itself already uses (e.g. `grassmann.jl`). +infer_rt(f, types::Tuple) = @static if isdefined(Base, :infer_return_type) + Base.infer_return_type(f, types) +else + Core.Compiler.return_type(f, Tuple{types...}) +end + +@testset "device allocator selection" begin + L = 6 + V = ℂ^2 + ψ_cpu = FiniteMPS(randn, ComplexF64, L, V, ℂ^4) + ψ = adapt(ArrType, ψ_cpu) + + M_cpu = TensorKit.storagetype(ψ_cpu) + M_dev = TensorKit.storagetype(ψ) + @test M_dev !== M_cpu + + # the whole point: a device-backed state is never handed a host allocator, whose intermediates + # would live in host memory - it allocates through its own storage type instead + for scheduler in (SerialScheduler(), DynamicScheduler()) + @test default_allocator(M_dev, scheduler) isa DefaultAllocator + @test !(default_allocator(M_dev, scheduler) isa Union{BufferAllocator, ManualAllocator}) + # and the state dispatches the same way as its storage type + @test typeof(default_allocator(ψ, scheduler)) === typeof(default_allocator(M_dev, scheduler)) + end + @test default_allocator(M_cpu, SerialScheduler()) isa BufferAllocator + @test default_allocator(M_cpu, DynamicScheduler()) isa ManualAllocator + + # Inference must pin the allocator, or every contraction downstream becomes a dynamic call. + @test isconcretetype(infer_rt(default_allocator, (typeof(ψ), SerialScheduler))) + @test infer_rt(default_allocator, (typeof(ψ), typeof(DynamicScheduler()))) === DefaultAllocator +end + +@testset "find_groundstate on device" begin + L = 8 + V = ℂ^2 + X = TensorMap(ComplexF64[0 1; 1 0], V ← V) + Z = TensorMap(ComplexF64[1 0; 0 -1], V ← V) + H_cpu = FiniteMPOHamiltonian( + fill(V, L), + (i => -2.0 * X for i in 1:L)..., + ((i, i + 1) => -(Z ⊗ Z) for i in 1:(L - 1))... + ) + + Random.seed!(0x51de) + ψ₀_cpu = FiniteMPS(randn, ComplexF64, L, V, ℂ^16) + H = adapt(ArrType, H_cpu) + ψ₀ = adapt(ArrType, ψ₀_cpu) + + alg = DMRG(; verbosity = 0, maxiter = 8) + ψ_ref, = find_groundstate(ψ₀_cpu, H_cpu, alg) + E_ref = expectation_value(ψ_ref, H_cpu) + + ψ, = find_groundstate(ψ₀, H, alg) + E = expectation_value(ψ, H) + + # the state must not have migrated off the device on the way through + @test TensorKit.storagetype(ψ) === TensorKit.storagetype(ψ₀) + @test real(E) ≈ real(E_ref) atol = 1.0e-6 +end + +@testset "timestep on device" begin + L = 6 + V = ℂ^2 + X = TensorMap(ComplexF64[0 1; 1 0], V ← V) + Z = TensorMap(ComplexF64[1 0; 0 -1], V ← V) + H_cpu = FiniteMPOHamiltonian( + fill(V, L), + (i => -1.0 * X for i in 1:L)..., + ((i, i + 1) => -(Z ⊗ Z) for i in 1:(L - 1))... + ) + + Random.seed!(0x7d1e) + ψ₀_cpu = complex(FiniteMPS(randn, ComplexF64, L, V, ℂ^8)) + H = adapt(ArrType, H_cpu) + ψ₀ = adapt(ArrType, ψ₀_cpu) + + for alg in (TDVP(), BUG()) + ψ_ref, = timestep(ψ₀_cpu, H_cpu, 0.0, 0.01, alg) + ψ, = timestep(ψ₀, H, 0.0, 0.01, alg) + + @test TensorKit.storagetype(ψ) === TensorKit.storagetype(ψ₀) + @test norm(ψ) ≈ norm(ψ₀) atol = 1.0e-6 + @test real(expectation_value(ψ, H)) ≈ real(expectation_value(ψ_ref, H_cpu)) atol = 1.0e-6 + end +end