Implement multi-threading using OhMyThreads and make it differentiable

Project.toml

@@ -13,6 +13,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

src/PEPSKit.jl

@@ -11,8 +11,10 @@ using LoggingExtras
 using MPSKit: loginit!, logiter!, logfinish!, logcancel!
 using MPSKitModels
 using FiniteDifferences
+using OhMyThreads
 
 include("utility/util.jl")
+include("utility/diffable_threads.jl")
 include("utility/svd.jl")
 include("utility/rotations.jl")
 include("utility/diffset.jl")
@@ -51,22 +53,45 @@ include("utility/symmetrization.jl")
     module Defaults
         const ctmrg_maxiter = 100
         const ctmrg_miniter = 4
-        const ctmrg_tol = 1e-12
-        const fpgrad_maxiter = 100
+        const ctmrg_tol = 1e-8
+        const fpgrad_maxiter = 30
         const fpgrad_tol = 1e-6
+        const ctmrgscheme = :simultaneous
+        const reuse_env = true
+        const trscheme = FixedSpaceTruncation()
+        const fwd_alg = TensorKit.SVD()
+        const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
+        const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
+        const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
+        const gradient_linsolver = KrylovKit.BiCGStab(;
+            maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
+        )
+        const iterscheme = :fixed
+        const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+        const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler())
     end
 
 Module containing default values that represent typical algorithm parameters.
 
-- `ctmrg_maxiter = 100`: Maximal number of CTMRG iterations per run
-- `ctmrg_miniter = 4`: Minimal number of CTMRG carried out
-- `ctmrg_tol = 1e-12`: Tolerance checking singular value and norm convergence
-- `fpgrad_maxiter = 100`: Maximal number of iterations for computing the CTMRG fixed-point gradient
-- `fpgrad_tol = 1e-6`: Convergence tolerance for the fixed-point gradient iteration
-- `rrule_alg = Arnoldi(; tol=1e1ctmrg_tol, krylovdim=48, verbosity=-1)`: Default cotangent linear problem algorithm
+- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run
+- `ctmrg_miniter`: Minimal number of CTMRG carried out
+- `ctmrg_tol`: Tolerance checking singular value and norm convergence
+- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient
+- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration
+- `ctmrgscheme`: Scheme for growing, projecting and renormalizing CTMRG environments
+- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment
+- `trscheme`: Truncation scheme for SVDs and other decompositions
+- `fwd_alg`: SVD algorithm that is used in the forward pass
+- `rrule_alg`: Reverse-rule for differentiating that SVD
+- `svd_alg`: Combination of `fwd_alg` and `rrule_alg`
+- `optimizer`: Optimization algorithm for PEPS ground-state optimization
+- `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm
+- `iterscheme`: Scheme for differentiating one CTMRG iteration
+- `gradient_alg`: Algorithm to compute the gradient fixed-point
+- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!`
 """
 module Defaults
-    using TensorKit, KrylovKit, OptimKit
+    using TensorKit, KrylovKit, OptimKit, OhMyThreads
     using PEPSKit: LinSolver, FixedSpaceTruncation, SVDAdjoint
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
@@ -77,17 +102,33 @@ module Defaults
     const sparse = false
     const reuse_env = true
     const trscheme = FixedSpaceTruncation()
-    const iterscheme = :fixed
     const fwd_alg = TensorKit.SVD()
     const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
     const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
     const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
     const gradient_linsolver = KrylovKit.BiCGStab(;
         maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
     )
+    const iterscheme = :fixed
     const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+
+    # OhMyThreads scheduler defaults
+    const scheduler = Ref{Scheduler}()
+    function set_scheduler!()
+        scheduler[] = Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler()
+        return nothing
+    end
+    function set_scheduler!(sc=OhMyThreads.Implementation.NotGiven(); kwargs...)
+        scheduler[] = OhMyThreads.Implementation._scheduler_from_userinput(sc; kwargs...)
+        return nothing
+    end
+
+    function __init__()
+        return set_scheduler!()
+    end
 end
 
+export set_scheduler!
 export SVDAdjoint, IterSVD, NonTruncSVDAdjoint
 export FixedSpaceTruncation, ProjectorAlg, CTMRG, CTMRGEnv, correlation_length
 export LocalOperator

src/algorithms/ctmrg/ctmrg.jl

@@ -8,13 +8,14 @@ have different spaces, this truncation style is different from `TruncationSpace`
 struct FixedSpaceTruncation <: TensorKit.TruncationScheme end
 
 """
-    struct ProjectorAlg{S}(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, verbosity=0)
-
-Algorithm struct collecting all projector related parameters.
-
-The `svd_alg` sets the SVD algorithm for decomposing the CTM environment. The truncation scheme
-has to be a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions
-on what kind of truncation scheme can be used.
+    struct ProjectorAlg{S}(; svd_alg=TensorKit.SVD(), trscheme=TensorKit.notrunc(),
+                           fixedspace=false, verbosity=0)
+
+Algorithm struct collecting all projector related parameters. The truncation scheme has to be
+a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions on what
+kind of truncation scheme can be used. If `fixedspace` is true, the truncation scheme is set to
+`truncspace(V)` where `V` is the environment bond space, adjusted to the corresponding
+environment direction/unit cell entry.
 """
 @kwdef struct ProjectorAlg{S<:SVDAdjoint,T}
     svd_alg::S = Defaults.svd_alg
@@ -179,34 +180,21 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N)
 
 """
     ctmrg_expand(state, envs, alg::CTMRG{M})
-    ctmrg_expand(dirs, state, envs::CTMRGEnv)
 
 Expand the environment by absorbing a new PEPS tensor.
 There are two modes of expansion: `M = :sequential` and `M = :simultaneous`.
 The first mode expands the environment in one direction at a time, for convenience towards
 the left. The second mode expands the environment in all four directions simultaneously.
-Alternatively, one can provide directly the `dirs` in which the environment is grown.
 """
 function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG)
     return ctmrg_expand([4, 1], state, envs)
 end
 function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG)
     return ctmrg_expand(1:4, state, envs)
 end
-# function ctmrg_expand(dirs, state, envs::CTMRGEnv)  # TODO: This doesn't AD due to length(::Nothing)...
-#     drc_combinations = collect(Iterators.product(dirs, axes(state)...))
-#     return map(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), drc_combinations)
-# end
-function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T}
-    Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C))
-    Q = Zygote.Buffer(Array{Qtype,3}(undef, length(dirs), size(state)...))
-    dirs_enum = [(i, dir) for (i, dir) in enumerate(dirs)]
-    drc_combinations = collect(Iterators.product(dirs_enum, axes(state)...))
-    @fwdthreads for (d, r, c) in drc_combinations
-        ec = EnlargedCorner(state, envs, (d[2], r, c))
-        Q[d[1], r, c] = TensorMap(ec, d[2])  # Explicitly construct EnlargedCorner for now
-    end
-    return copy(Q)
+function ctmrg_expand(dirs, state, envs::CTMRGEnv)
+    coordinates = eachcoordinate(state, dirs)
+    return dtmap(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), coordinates)
 end
 
 # ======================================================================================== #
@@ -223,15 +211,10 @@ function ctmrg_projectors(
     enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
 ) where {C,E}
     projector_alg = alg.projector_alg
-    # pre-allocation
-    Prtype = tensormaptype(spacetype(E), numin(E), numout(E), storagetype(E))
-    P_bottom = Zygote.Buffer(envs.edges, axes(envs.corners, 2), axes(envs.corners, 3))
-    P_top = Zygote.Buffer(envs.edges, Prtype, axes(envs.corners, 2), axes(envs.corners, 3))
     ϵ = zero(real(scalartype(envs)))
 
-    directions = collect(Iterators.product(axes(envs.corners, 2), axes(envs.corners, 3)))
-    # @fwdthreads for (r, c) in directions
-    for (r, c) in directions
+    coordinates = eachcoordinate(envs)
+    projectors = dtmap(coordinates) do (r, c)
         # SVD half-infinite environment
         r′ = _prev(r, size(envs.corners, 2))
         QQ = halfinfinite_environment(enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
@@ -250,26 +233,23 @@ function ctmrg_projectors(
         end
 
         # Compute projectors
-        P_bottom[r, c], P_top[r, c] = build_projectors(
-            U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c]
-        )
+        return build_projectors(U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
     end
 
-    return (copy(P_bottom), copy(P_top)), (; err=ϵ)
+    return (map(first, projectors), map(last, projectors)), (; err=ϵ)
 end
 function ctmrg_projectors(
     enlarged_envs, envs::CTMRGEnv{C,E}, alg::SimultaneousCTMRG
 ) where {C,E}
     projector_alg = alg.projector_alg
     # pre-allocation
-    P_left, P_right = Zygote.Buffer.(projector_type(envs.edges))
     U, V = Zygote.Buffer.(projector_type(envs.edges))
     # Corner type but with real numbers
     S = Zygote.Buffer(U.data, tensormaptype(spacetype(C), 1, 1, real(scalartype(E))))
 
     ϵ = zero(real(scalartype(envs)))
-    drc_combinations = collect(Iterators.product(axes(envs.corners)...))
-    @fwdthreads for (dir, r, c) in drc_combinations
+    coordinates = eachcoordinate(envs, 1:4)
+    projectors = dtmap(coordinates) do (dir, r, c)
         # Row-column index of next enlarged corner
         next_rc = if dir == 1
             (r, _next(c, size(envs.corners, 3)))
@@ -303,7 +283,7 @@ function ctmrg_projectors(
         end
 
         # Compute projectors
-        P_left[dir, r, c], P_right[dir, r, c] = build_projectors(
+        return build_projectors(
             U_local,
             S_local,
             V_local,
@@ -312,7 +292,9 @@ function ctmrg_projectors(
         )
     end
 
-    return (copy(P_left), copy(P_right)), (; err=ϵ, U=copy(U), S=copy(S), V=copy(V))
+    P_left = map(first, projectors)
+    P_right = map(last, projectors)
+    return (P_left, P_right), (; err=ϵ, U=copy(U), S=copy(S), V=copy(V))
 end
 
 """
@@ -376,9 +358,7 @@ function ctmrg_renormalize(projectors, state, envs, ::SequentialCTMRG)
     end
 
     # Apply projectors to renormalize corners and edges
-    coordinates = collect(Iterators.product(axes(state)...))
-    # @fwdthreads for (r, c) in coordinates
-    for (r, c) in coordinates
+    for (r, c) in eachcoordinate(state)
         C_southwest = renormalize_bottom_corner((r, c), envs, projectors)
         corners[SOUTHWEST, r, c] = C_southwest / norm(C_southwest)
 
@@ -392,12 +372,9 @@ function ctmrg_renormalize(projectors, state, envs, ::SequentialCTMRG)
     return CTMRGEnv(copy(corners), copy(edges))
 end
 function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::SimultaneousCTMRG)
-    corners = Zygote.Buffer(envs.corners)
-    edges = Zygote.Buffer(envs.edges)
     P_left, P_right = projectors
-
-    drc_combinations = collect(Iterators.product(axes(envs.corners)...))
-    @fwdthreads for (dir, r, c) in drc_combinations
+    coordinates = eachcoordinate(envs, 1:4)
+    corners_edges = dtmap(coordinates) do (dir, r, c)
         if dir == NORTH
             corner = renormalize_northwest_corner((r, c), enlarged_envs, P_left, P_right)
             edge = renormalize_north_edge((r, c), envs, P_left, P_right, state)
@@ -411,9 +388,8 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::Simultaneou
             corner = renormalize_southwest_corner((r, c), enlarged_envs, P_left, P_right)
             edge = renormalize_west_edge((r, c), envs, P_left, P_right, state)
         end
-        corners[dir, r, c] = corner / norm(corner)
-        edges[dir, r, c] = edge / norm(edge)
+        return corner / norm(corner), edge / norm(edge)
     end
 
-    return CTMRGEnv(copy(corners), copy(edges))
+    return CTMRGEnv(map(first, corners_edges), map(last, corners_edges))
 end

src/algorithms/ctmrg/gaugefix.jl

@@ -8,14 +8,14 @@ element-wise converged to `envprev`.
 """
 function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
     # Check if spaces in envprev and envfinal are the same
-    same_spaces = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
+    same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) &&
             space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c])
     end
     @assert all(same_spaces) "Spaces of envprev and envfinal are not the same"
 
     # Try the "general" algorithm from https://arxiv.org/abs/2311.11894
-    signs = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
+    signs = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         # Gather edge tensors and pretend they're InfiniteMPSs
         if dir == NORTH
             Tsprev = circshift(envprev.edges[dir, r, :], 1 - c)
@@ -70,7 +70,7 @@ end
 
 # Explicit fixing of relative phases (doing this compactly in a loop is annoying)
 function fix_relative_phases(envfinal::CTMRGEnv, signs)
-    corners_fixed = map(Iterators.product(axes(envfinal.corners)...)) do (dir, r, c)
+    corners_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         if dir == NORTHWEST
             fix_gauge_northwest_corner((r, c), envfinal, signs)
         elseif dir == NORTHEAST
@@ -82,7 +82,7 @@ function fix_relative_phases(envfinal::CTMRGEnv, signs)
         end
     end
 
-    edges_fixed = map(Iterators.product(axes(envfinal.corners)...)) do (dir, r, c)
+    edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         if dir == NORTHWEST
             fix_gauge_north_edge((r, c), envfinal, signs)
         elseif dir == NORTHEAST
@@ -99,7 +99,8 @@ end
 function fix_relative_phases(
     U::Array{Ut,3}, V::Array{Vt,3}, signs
 ) where {Ut<:AbstractTensorMap,Vt<:AbstractTensorMap}
-    U_fixed = map(Iterators.product(axes(U)...)) do (dir, r, c)
+    U_fixed = map(CartesianIndices(U)) do I
+        dir, r, c = I.I
         if dir == NORTHWEST
             fix_gauge_north_left_vecs((r, c), U, signs)
         elseif dir == NORTHEAST
@@ -111,7 +112,8 @@ function fix_relative_phases(
         end
     end
 
-    V_fixed = map(Iterators.product(axes(V)...)) do (dir, r, c)
+    V_fixed = map(CartesianIndices(V)) do I
+        dir, r, c = I.I
         if dir == NORTHWEST
             fix_gauge_north_right_vecs((r, c), V, signs)
         elseif dir == NORTHEAST
@@ -191,7 +193,8 @@ function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol
     @debug "maxᵢⱼ|Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmax   mean |Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmean"
 
     # Check differences for all tensors in unit cell to debug properly
-    for (dir, r, c) in Iterators.product(axes(envfinal.edges)...)
+    for I in CartesianIndices(ΔT)
+        dir, r, c = I.I
         @debug(
             "$((dir, r, c)): all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ",
             all(x -> abs(x) < atol, convert(Array, ΔC[dir, r, c])),

src/algorithms/toolbox.jl

@@ -1,9 +1,10 @@
 function MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, envs::CTMRGEnv)
     checklattice(peps, O)
-    return sum(O.terms) do (inds, operator)  # TODO: parallelize this map, especially for the backwards pass
+    term_vals = dtmap([O.terms...]) do (inds, operator)  # OhMyThreads can't iterate over O.terms directly
         contract_localoperator(inds, operator, peps, peps, envs) /
         contract_localnorm(inds, peps, peps, envs)
     end
+    return sum(term_vals)
 end
 
 function costfun(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator)

src/environments/ctmrg_environments.jl

@@ -105,7 +105,8 @@ function CTMRGEnv(
     corners = Array{C_type}(undef, 4, size(Ds_north)...)
     edges = Array{T_type}(undef, 4, size(Ds_north)...)
 
-    for (r, c) in Iterators.product(axes(Ds_north)...)
+    for I in CartesianIndices(Ds_north)
+        r, c = I.I
         edges[NORTH, r, c] = _edge_tensor(
             f,
             T,
@@ -371,6 +372,13 @@ function Base.rotl90(env::CTMRGEnv{C,T}) where {C,T}
 end
 
 Base.eltype(env::CTMRGEnv) = eltype(env.corners[1])
+Base.axes(x::CTMRGEnv, args...) = axes(x.corners, args...)
+function eachcoordinate(x::CTMRGEnv)
+    return collect(Iterators.product(axes(x, 2), axes(x, 3)))
+end
+function eachcoordinate(x::CTMRGEnv, dirs)
+    return collect(Iterators.product(dirs, axes(x, 2), axes(x, 3)))
+end
 
 # In-place update of environment
 function update!(env::CTMRGEnv{C,T}, env´::CTMRGEnv{C,T}) where {C,T}

src/operators/infinitepepo.jl

@@ -125,7 +125,8 @@ function initializePEPS(
     T::InfinitePEPO{<:PEPOTensor{S}}, vspace::S
 ) where {S<:ElementarySpace}
     Pspaces = Array{S,2}(undef, size(T, 1), size(T, 2))
-    for (i, j) in product(1:size(T, 1), 1:size(T, 2))
+    for i in axes(T, 1)
+        j in axes(T, 2)
         Pspaces[i, j] = space(T, i, j)
     end
     Nspaces = repeat([vspace], size(T, 1), size(T, 2))

src/states/infinitepeps.jl

@@ -124,6 +124,12 @@ Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...))
 Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...)
 Base.setindex!(T::InfinitePEPS, args...) = (Base.setindex!(T.A, args...); T)
 Base.axes(T::InfinitePEPS, args...) = axes(T.A, args...)
+function eachcoordinate(x::InfinitePEPS)
+    return collect(Iterators.product(axes(x)...))
+end
+function eachcoordinate(x::InfinitePEPS, dirs)
+    return collect(Iterators.product(dirs, axes(x, 1), axes(x, 2)))
+end
 TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1)
 
 ## Math