[Perf] Another round of performance and type-stability improvements (#229)

* Optimize contraction: `renormalize_west_edge`
* Improve type stability of CTMRG
* Fix type stability for fixedspacetruncation
* Avoid boxed variables in closure
* patch svd algorithm when krylov dimension is too low

Author: lkdvos

Project.toml

@@ -43,6 +43,7 @@ Random = "1"
 Statistics = "1"
 TensorKit = "0.14.6"
 TensorOperations = "5"
+TestExtras = "0.3"
 VectorInterface = "0.4, 0.5"
 Zygote = "0.6, 0.7"
 julia = "1.10"
@@ -53,6 +54,7 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
 [targets]
-test = ["Test", "SafeTestsets", "ChainRulesTestUtils", "FiniteDifferences", "QuadGK"]
+test = ["Test", "TestExtras", "SafeTestsets", "ChainRulesTestUtils", "FiniteDifferences", "QuadGK"]

src/PEPSKit.jl

@@ -15,7 +15,7 @@ import MPSKit: leading_boundary, loginit!, logiter!, logfinish!, logcancel!, phy
 
 using MPSKitModels
 using FiniteDifferences
-using OhMyThreads: tmap
+using OhMyThreads: tmap, tmap!
 using DocStringExtensions
 
 include("Defaults.jl")  # Include first to allow for docstring interpolation with Defaults values

src/algorithms/contractions/ctmrg_contractions.jl

@@ -1310,12 +1310,19 @@ end
 function renormalize_west_edge(
     E_west::CTMRG_PEPS_EdgeTensor, P_bottom, P_top, A::PEPSSandwich
 )
-    return @autoopt @tensor edge[χ_S D_Eab D_Ebe; χ_N] :=
-        E_west[χ1 D1 D2; χ2] *
-        ket(A)[d; D3 D_Eab D5 D1] *
-        conj(bra(A)[d; D4 D_Ebe D6 D2]) *
-        P_bottom[χ2 D3 D4; χ_N] *
-        P_top[χ_S; χ1 D5 D6]
+    # starting with P_top to save one permute in the end
+    return @tensor begin
+        # already putting χE in front here to make next permute cheaper
+        PE[χS χNW DSb DWb; DSt DWt] := P_top[χS; χSW DSt DSb] * E_west[χSW DWt DWb; χNW]
+
+        PEket[χS χNW DNt DEt; DSb DWb d] :=
+            PE[χS χNW DSb DWb; DSt DWt] * ket(A)[d; DNt DEt DSt DWt]
+
+        corner[χS DEt DEb; χNW DNt DNb] :=
+            PEket[χS χNW DNt DEt; DSb DWb d] * conj(bra(A)[d; DNb DEb DSb DWb])
+
+        edge[χS DEt DEb; χN] := corner[χS DEt DEb; χNW DNt DNb] * P_bottom[χNW DNt DNb; χN]
+    end
 end
 function renormalize_west_edge(E_west::CTMRG_PF_EdgeTensor, P_bottom, P_top, A::PFTensor)
     return @autoopt @tensor edge[χ_S D_E; χ_N] :=

src/algorithms/correlators.jl

@@ -120,7 +120,7 @@ function end_correlator_numerator(
     C_southeast = env.corners[SOUTHEAST, _next(r, end), _next(c, end)]
     sandwich = (above[mod1(r, end), mod1(c, end)], below[mod1(r, end), mod1(c, end)])
 
-    return @autoopt @tensor contractcheck = true V[χSW DWt dstring DWb; χNW] *
+    return @autoopt @tensor V[χSW DWt dstring DWb; χNW] *
         E_south[χSSE DSt DSb; χSW] *
         E_east[χNEE DEt DEb; χSEE] *
         E_north[χNW DNt DNb; χNNE] *

src/algorithms/ctmrg/ctmrg.jl

@@ -110,14 +110,12 @@ end
 function leading_boundary(
     env₀::CTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm
 )
-    CS = map(x -> tsvd(x)[2], env₀.corners)
-    TS = map(x -> tsvd(x)[2], env₀.edges)
-
-    η = one(real(scalartype(network)))
-    env = deepcopy(env₀)
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
-
     return LoggingExtras.withlevel(; alg.verbosity) do
+        env = deepcopy(env₀)
+        CS = map(x -> tsvd(x)[2], env₀.corners)
+        TS = map(x -> tsvd(x)[2], env₀.edges)
+        η = one(real(scalartype(network)))
         ctmrg_loginit!(log, η, network, env₀)
         local info
         for iter in 1:(alg.maxiter)

src/algorithms/ctmrg/sequential.jl

@@ -24,12 +24,12 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a
 * `svd_alg::Union{<:SVDAdjoint,NamedTuple}`
 * `projector_alg::Symbol=:$(Defaults.projector_alg)`
 """
-struct SequentialCTMRG <: CTMRGAlgorithm
+struct SequentialCTMRG{P<:ProjectorAlgorithm} <: CTMRGAlgorithm
     tol::Float64
     maxiter::Int
     miniter::Int
     verbosity::Int
-    projector_alg::ProjectorAlgorithm
+    projector_alg::P
 end
 function SequentialCTMRG(; kwargs...)
     return CTMRGAlgorithm(; alg=:sequential, kwargs...)
@@ -81,14 +81,18 @@ for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequen
 """
 function sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm)
     coordinates = eachcoordinate(env)[:, col]
-    proj_and_info = dtmap(coordinates) do (r, c)
+    T_dst = Base.promote_op(
+        sequential_projectors, NTuple{3,Int}, typeof(network), typeof(env), typeof(alg)
+    )
+    proj_and_info = similar(coordinates, T_dst)
+    proj_and_info′::typeof(proj_and_info) = dtmap!!(proj_and_info, coordinates) do (r, c)
         trscheme = truncation_scheme(alg, env.edges[WEST, _prev(r, size(env, 2)), c])
         proj, info = sequential_projectors(
             (WEST, r, c), network, env, @set(alg.trscheme = trscheme)
         )
         return proj, info
     end
-    return _split_proj_and_info(proj_and_info)
+    return _split_proj_and_info(proj_and_info′)
 end
 function sequential_projectors(
     coordinate::NTuple{3,Int}, network, env::CTMRGEnv, alg::HalfInfiniteProjector

src/algorithms/ctmrg/simultaneous.jl

@@ -23,12 +23,12 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a
 * `svd_alg::Union{<:SVDAdjoint,NamedTuple}`
 * `projector_alg::Symbol=:$(Defaults.projector_alg)`
 """
-struct SimultaneousCTMRG <: CTMRGAlgorithm
+struct SimultaneousCTMRG{P<:ProjectorAlgorithm} <: CTMRGAlgorithm
     tol::Float64
     maxiter::Int
     miniter::Int
     verbosity::Int
-    projector_alg::ProjectorAlgorithm
+    projector_alg::P
 end
 function SimultaneousCTMRG(; kwargs...)
     return CTMRGAlgorithm(; alg=:simultaneous, kwargs...)
@@ -37,9 +37,15 @@ end
 CTMRG_SYMBOLS[:simultaneous] = SimultaneousCTMRG
 
 function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG)
-    enlarged_corners = dtmap(eachcoordinate(network, 1:4)) do idx
-        return TensorMap(EnlargedCorner(network, env, idx))
-    end  # expand environment
+    coordinates = eachcoordinate(network, 1:4)
+    T_corners = Base.promote_op(
+        TensorMap ∘ EnlargedCorner, typeof(network), typeof(env), eltype(coordinates)
+    )
+    enlarged_corners′ = similar(coordinates, T_corners)
+    enlarged_corners::typeof(enlarged_corners′) =
+        dtmap!!(enlarged_corners′, eachcoordinate(network, 1:4)) do idx
+            return TensorMap(EnlargedCorner(network, env, idx))
+        end  # expand environment
     projectors, info = simultaneous_projectors(enlarged_corners, env, alg.projector_alg)  # compute projectors on all coordinates
     env′ = renormalize_simultaneously(enlarged_corners, projectors, network, env)  # renormalize enlarged corners
     return env′, info
@@ -72,25 +78,36 @@ enlarged corners or on a specific `coordinate`.
 function simultaneous_projectors(
     enlarged_corners::Array{E,3}, env::CTMRGEnv, alg::ProjectorAlgorithm
 ) where {E}
-    proj_and_info = dtmap(eachcoordinate(env, 1:4)) do coordinate
-        coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...)
-        trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...])
-        return simultaneous_projectors(
-            coordinate, enlarged_corners, @set(alg.trscheme = trscheme)
-        )
-    end
+    coordinates = eachcoordinate(env, 1:4)
+    T_dst = Base.promote_op(
+        simultaneous_projectors,
+        NTuple{3,Int},
+        typeof(enlarged_corners),
+        typeof(env),
+        typeof(alg),
+    )
+    proj_and_info′ = similar(coordinates, T_dst)
+    proj_and_info::typeof(proj_and_info′) =
+        dtmap!!(proj_and_info′, coordinates) do coordinate
+            return simultaneous_projectors(coordinate, enlarged_corners, env, alg)
+        end
     return _split_proj_and_info(proj_and_info)
 end
 function simultaneous_projectors(
-    coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector
+    coordinate, enlarged_corners::Array{E,3}, env, alg::HalfInfiniteProjector
 ) where {E}
-    coordinate′ = _next_coordinate(coordinate, size(enlarged_corners)[2:3]...)
+    coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...)
+    trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...])
+    alg′ = @set alg.trscheme = trscheme
     ec = (enlarged_corners[coordinate...], enlarged_corners[coordinate′...])
-    return compute_projector(ec, coordinate, alg)
+    return compute_projector(ec, coordinate, alg′)
 end
 function simultaneous_projectors(
-    coordinate, enlarged_corners::Array{E,3}, alg::FullInfiniteProjector
+    coordinate, enlarged_corners::Array{E,3}, env, alg::FullInfiniteProjector
 ) where {E}
+    coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...)
+    trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...])
+    alg′ = @set alg.trscheme = trscheme
     rowsize, colsize = size(enlarged_corners)[2:3]
     coordinate2 = _next_coordinate(coordinate, rowsize, colsize)
     coordinate3 = _next_coordinate(coordinate2, rowsize, colsize)
@@ -101,7 +118,7 @@ function simultaneous_projectors(
         enlarged_corners[coordinate2...],
         enlarged_corners[coordinate3...],
     )
-    return compute_projector(ec, coordinate, alg)
+    return compute_projector(ec, coordinate, alg′)
 end
 
 """
@@ -112,22 +129,33 @@ Renormalize all enlarged corners and edges simultaneously.
 function renormalize_simultaneously(enlarged_corners, projectors, network, env)
     P_left, P_right = projectors
     coordinates = eachcoordinate(env, 1:4)
-    corners_edges = dtmap(coordinates) do (dir, r, c)
-        if dir == NORTH
-            corner = renormalize_northwest_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_north_edge((r, c), env, P_left, P_right, network)
-        elseif dir == EAST
-            corner = renormalize_northeast_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_east_edge((r, c), env, P_left, P_right, network)
-        elseif dir == SOUTH
-            corner = renormalize_southeast_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_south_edge((r, c), env, P_left, P_right, network)
-        elseif dir == WEST
-            corner = renormalize_southwest_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_west_edge((r, c), env, P_left, P_right, network)
+    T_CE = Tuple{cornertype(env),edgetype(env)}
+    corners_edges′ = similar(coordinates, T_CE)
+    corners_edges::typeof(corners_edges′) =
+        dtmap!!(corners_edges′, coordinates) do (dir, r, c)
+            if dir == NORTH
+                corner = renormalize_northwest_corner(
+                    (r, c), enlarged_corners, P_left, P_right
+                )
+                edge = renormalize_north_edge((r, c), env, P_left, P_right, network)
+            elseif dir == EAST
+                corner = renormalize_northeast_corner(
+                    (r, c), enlarged_corners, P_left, P_right
+                )
+                edge = renormalize_east_edge((r, c), env, P_left, P_right, network)
+            elseif dir == SOUTH
+                corner = renormalize_southeast_corner(
+                    (r, c), enlarged_corners, P_left, P_right
+                )
+                edge = renormalize_south_edge((r, c), env, P_left, P_right, network)
+            elseif dir == WEST
+                corner = renormalize_southwest_corner(
+                    (r, c), enlarged_corners, P_left, P_right
+                )
+                edge = renormalize_west_edge((r, c), env, P_left, P_right, network)
+            end
+            return corner / norm(corner), edge / norm(edge)
         end
-        return corner / norm(corner), edge / norm(edge)
-    end
 
     return CTMRGEnv(map(first, corners_edges), map(last, corners_edges))
 end

src/algorithms/ctmrg/sparse_environments.jl

@@ -56,6 +56,8 @@ function EnlargedCorner(network::InfiniteSquareNetwork, env, coordinates)
             network[r, c],
             dir,
         )
+    else
+        throw(ArgumentError(lazy"Invalid direction $dir"))
     end
 end
 
@@ -73,6 +75,8 @@ function TensorKit.TensorMap(Q::EnlargedCorner)
         return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.A)
     elseif Q.dir == SOUTHWEST
         return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.A)
+    else
+        throw(ArgumentError(lazy"Invalid direction $dir"))
     end
 end
 

src/utility/diffable_threads.jl

@@ -8,6 +8,8 @@ All calls of `dtmap` inside of PEPSKit use the threading scheduler stored inside
 """
 dtmap(args...; scheduler=Defaults.scheduler[]) = tmap(args...; scheduler)
 
+dtmap!!(args...; scheduler=Defaults.scheduler[]) = tmap!(args...; scheduler)
+
 # Follows the `map` rrule from ChainRules.jl but specified for the case of one AbstractArray that is being mapped
 # https://github.com/JuliaDiff/ChainRules.jl/blob/e245d50a1ae56ce46fc8c1f0fe9b925964f1146e/src/rulesets/Base/base.jl#L243
 function ChainRulesCore.rrule(
@@ -33,6 +35,22 @@ function ChainRulesCore.rrule(
     return y, dtmap_pullback
 end
 
+function ChainRulesCore.rrule(
+    config::RuleConfig{>:HasReverseMode},
+    ::typeof(dtmap!!),
+    f,
+    C′::AbstractArray,
+    A::AbstractArray;
+    kwargs...,
+)
+    C, dtmap_pullback = rrule(config, dtmap, f, A; kwargs...)
+    function dtmap!!_pullback(dy)
+        dtmap, df, dA = dtmap_pullback(dy)
+        return dtmap, df, NoTangent, dA
+    end
+    return C, dtmap!!_pullback
+end
+
 """
     @fwdthreads(ex)
 

src/utility/svd.jl

@@ -290,7 +290,11 @@ function TensorKit._compute_svddata!(
             V = V[1:howmany, :]
         else
             x₀ = alg.start_vector(b)
-            S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg)
+            svd_alg = alg.alg
+            if howmany > alg.alg.krylovdim
+                svd_alg = @set svd_alg.krylovdim = round(Int, howmany * 1.2)
+            end
+            S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, svd_alg)
             if info.converged < howmany  # Fall back to dense SVD if not properly converged
                 @warn "Iterative SVD did not converge for block $c, falling back to dense SVD"
                 U, S, V = TensorKit.MatrixAlgebra.svd!(b, TensorKit.SDD())