[Perf] EnlargedCorner optimizations (#214)

* Store `dir` in `EnlargedCorner`

* Manually fix enlarged corner contractions

* Update contractions

---

Co-authored-by: Olivier Gauthe
<olivier.gauthe.2011+github@polytechnique.org>

Author: lkdvos

src/algorithms/contractions/ctmrg_contractions.jl

@@ -6,6 +6,20 @@ const CTMRGCornerTensor{T,S} = AbstractTensorMap{T,S,1,1}
 # Enlarged corner contractions
 # ----------------------------
 
+#=
+These contractions are hand-optimized by the following heuristics:
+
+1. ensure contraction order gives minimal scaling in χ = D²
+2. ensure dominant permutation is as efficient as possible by making large legs contiguous,
+    ie moving them to the front
+
+This second part is mostly important for dealing with non-abelian symmetries, where the
+    permutations are strongly non-negligable.
+
+For a small benchmark study:
+https://gist.github.com/lkdvos/a562c2b09ef461398729ccefdab34745
+=#
+
 """
 $(SIGNATURES)
 
@@ -25,12 +39,18 @@ function enlarge_northwest_corner(
     E_north::CTMRG_PEPS_EdgeTensor,
     A::PEPSSandwich,
 )
-    return @autoopt @tensor corner[χ_S D_Sabove D_Sbelow; χ_E D_Eabove D_Ebelow] :=
-        E_west[χ_S D1 D2; χ1] *
-        C_northwest[χ1; χ2] *
-        E_north[χ2 D3 D4; χ_E] *
-        ket(A)[d; D3 D_Eabove D_Sabove D1] *
-        conj(bra(A)[d; D4 D_Ebelow D_Sbelow D2])
+    return @tensor begin
+        EC[χS DWt DWb; χ2] := E_west[χS DWt DWb; χ1] * C_northwest[χ1; χ2]
+
+        # already putting χE in front here to make next permute cheaper
+        ECE[χS χE DWb DNb; DWt DNt] := EC[χS DWt DWb; χ2] * E_north[χ2 DNt DNb; χE]
+
+        ECEket[χS χE DEt DSt; DWb DNb d] :=
+            ECE[χS χE DWb DNb; DWt DNt] * ket(A)[d; DNt DEt DSt DWt]
+
+        corner[χS DSt DSb; χE DEt DEb] :=
+            ECEket[χS χE DEt DSt; DWb DNb d] * conj(bra(A)[d; DNb DEb DSb DWb])
+    end
 end
 function enlarge_northwest_corner(
     E_west::CTMRG_PF_EdgeTensor,
@@ -61,12 +81,18 @@ function enlarge_northeast_corner(
     E_east::CTMRG_PEPS_EdgeTensor,
     A::PEPSSandwich,
 )
-    return @autoopt @tensor corner[χ_W D_Wabove D_Wbelow; χ_S D_Sabove D_Sbelow] :=
-        E_north[χ_W D1 D2; χ1] *
-        C_northeast[χ1; χ2] *
-        E_east[χ2 D3 D4; χ_S] *
-        ket(A)[d; D1 D3 D_Sabove D_Wabove] *
-        conj(bra(A)[d; D2 D4 D_Sbelow D_Wbelow])
+    return @tensor begin
+        EC[χW DNt DNb; χ2] := E_north[χW DNt DNb; χ1] * C_northeast[χ1; χ2]
+
+        # already putting χE in front here to make next permute cheaper
+        ECE[χW χS DNb DEb; DNt DEt] := EC[χW DNt DNb; χ2] * E_east[χ2 DEt DEb; χS]
+
+        ECEket[χW χS DSt DWt; DNb DEb d] :=
+            ECE[χW χS DNb DEb; DNt DEt] * ket(A)[d; DNt DEt DSt DWt]
+
+        corner[χW DWt DWb; χS DSt DSb] :=
+            ECEket[χW χS DSt DWt; DNb DEb d] * conj(bra(A)[d; DNb DEb DSb DWb])
+    end
 end
 function enlarge_northeast_corner(
     E_north::CTMRG_PF_EdgeTensor,
@@ -97,12 +123,18 @@ function enlarge_southeast_corner(
     E_south::CTMRG_PEPS_EdgeTensor,
     A::PEPSSandwich,
 )
-    return @autoopt @tensor corner[χ_N D_Nabove D_Nbelow; χ_W D_Wabove D_Wbelow] :=
-        E_east[χ_N D1 D2; χ1] *
-        C_southeast[χ1; χ2] *
-        E_south[χ2 D3 D4; χ_W] *
-        ket(A)[d; D_Nabove D1 D3 D_Wabove] *
-        conj(bra(A)[d; D_Nbelow D2 D4 D_Wbelow])
+    return @tensor begin
+        EC[χN DEt DEb; χ2] := E_east[χN DEt DEb; χ1] * C_southeast[χ1; χ2]
+
+        # already putting χE in front here to make next permute cheaper
+        ECE[χN χW DEb DSb; DEt DSt] := EC[χN DEt DEb; χ2] * E_south[χ2 DSt DSb; χW]
+
+        ECEket[χN χW DNt DWt; DEb DSb d] :=
+            ECE[χN χW DEb DSb; DEt DSt] * ket(A)[d; DNt DEt DSt DWt]
+
+        corner[χN DNt DNb; χW DWt DWb] :=
+            ECEket[χN χW DNt DWt; DEb DSb d] * conj(bra(A)[d; DNb DEb DSb DWb])
+    end
 end
 function enlarge_southeast_corner(
     E_east::CTMRG_PF_EdgeTensor,
@@ -133,12 +165,18 @@ function enlarge_southwest_corner(
     E_west::CTMRG_PEPS_EdgeTensor,
     A::PEPSSandwich,
 )
-    return @autoopt @tensor corner[χ_E D_Eabove D_Ebelow; χ_N D_Nabove D_Nbelow] :=
-        E_south[χ_E D1 D2; χ1] *
-        C_southwest[χ1; χ2] *
-        E_west[χ2 D3 D4; χ_N] *
-        ket(A)[d; D_Nabove D_Eabove D1 D3] *
-        conj(bra(A)[d; D_Nbelow D_Ebelow D2 D4])
+    return @tensor begin
+        EC[χE DSt DSb; χ2] := E_south[χE DSt DSb; χ1] * C_southwest[χ1; χ2]
+
+        # already putting χE in front here to make next permute cheaper
+        ECE[χE χN DSb DWb; DSt DWt] := EC[χE DSt DSb; χ2] * E_west[χ2 DWt DWb; χN]
+
+        ECEket[χE χN DNt DEt; DSb DWb d] :=
+            ECE[χE χN DSb DWb; DSt DWt] * ket(A)[d; DNt DEt DSt DWt]
+
+        corner[χE DEt DEb; χN DNt DNb] :=
+            ECEket[χE χN DNt DEt; DSb DWb d] * conj(bra(A)[d; DNb DEb DSb DWb])
+    end
 end
 function enlarge_southwest_corner(
     E_south::CTMRG_PF_EdgeTensor,

src/algorithms/ctmrg/sequential.jl

@@ -95,8 +95,8 @@ function sequential_projectors(
 )
     _, r, c = coordinate
     r′ = _prev(r, size(env, 2))
-    Q1 = TensorMap(EnlargedCorner(network, env, (SOUTHWEST, r, c)), SOUTHWEST)
-    Q2 = TensorMap(EnlargedCorner(network, env, (NORTHWEST, r′, c)), NORTHWEST)
+    Q1 = TensorMap(EnlargedCorner(network, env, (SOUTHWEST, r, c)))
+    Q2 = TensorMap(EnlargedCorner(network, env, (NORTHWEST, r′, c)))
     return compute_projector((Q1, Q2), coordinate, alg)
 end
 function sequential_projectors(
@@ -107,10 +107,10 @@ function sequential_projectors(
     coordinate_ne = _next_coordinate(coordinate_nw, rowsize, colsize)
     coordinate_se = _next_coordinate(coordinate_ne, rowsize, colsize)
     ec = (
-        TensorMap(EnlargedCorner(network, env, coordinate_se), SOUTHEAST),
-        TensorMap(EnlargedCorner(network, env, coordinate), SOUTHWEST),
-        TensorMap(EnlargedCorner(network, env, coordinate_nw), NORTHWEST),
-        TensorMap(EnlargedCorner(network, env, coordinate_ne), NORTHEAST),
+        TensorMap(EnlargedCorner(network, env, coordinate_se)),
+        TensorMap(EnlargedCorner(network, env, coordinate)),
+        TensorMap(EnlargedCorner(network, env, coordinate_nw)),
+        TensorMap(EnlargedCorner(network, env, coordinate_ne)),
     )
     return compute_projector(ec, coordinate, alg)
 end

src/algorithms/ctmrg/simultaneous.jl

@@ -38,7 +38,7 @@ 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), idx[1])
+        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

src/algorithms/ctmrg/sparse_environments.jl

@@ -20,6 +20,7 @@ struct EnlargedCorner{TC,TE,TA}
     E_1::TE
     E_2::TE
     A::TA
+    dir::Int
 end
 function EnlargedCorner(network::InfiniteSquareNetwork, env, coordinates)
     dir, r, c = coordinates
@@ -29,45 +30,48 @@ function EnlargedCorner(network::InfiniteSquareNetwork, env, coordinates)
             env.edges[WEST, r, _prev(c, end)],
             env.edges[NORTH, _prev(r, end), c],
             network[r, c],
+            dir,
         )
     elseif dir == NORTHEAST
         return EnlargedCorner(
             env.corners[NORTHEAST, _prev(r, end), _next(c, end)],
             env.edges[NORTH, _prev(r, end), c],
             env.edges[EAST, r, _next(c, end)],
             network[r, c],
+            dir,
         )
     elseif dir == SOUTHEAST
         return EnlargedCorner(
             env.corners[SOUTHEAST, _next(r, end), _next(c, end)],
             env.edges[EAST, r, _next(c, end)],
             env.edges[SOUTH, _next(r, end), c],
             network[r, c],
+            dir,
         )
     elseif dir == SOUTHWEST
         return EnlargedCorner(
             env.corners[SOUTHWEST, _next(r, end), _prev(c, end)],
             env.edges[SOUTH, _next(r, end), c],
             env.edges[WEST, r, _prev(c, end)],
             network[r, c],
+            dir,
         )
     end
 end
 
 """
-    TensorMap(Q::EnlargedCorner, dir::Int)
+    TensorMap(Q::EnlargedCorner)
 
-Instantiate enlarged corner as `TensorMap` where `dir` selects the correct contraction
-direction, i.e. the way the environment and PEPS tensors connect.
+Instantiate enlarged corner as a `TensorMap`.
 """
-function TensorKit.TensorMap(Q::EnlargedCorner, dir::Int)
-    if dir == NORTHWEST
+function TensorKit.TensorMap(Q::EnlargedCorner)
+    if Q.dir == NORTHWEST
         return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.A)
-    elseif dir == NORTHEAST
+    elseif Q.dir == NORTHEAST
         return enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.A)
-    elseif dir == SOUTHEAST
+    elseif Q.dir == SOUTHEAST
         return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.A)
-    elseif dir == SOUTHWEST
+    elseif Q.dir == SOUTHWEST
         return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.A)
     end
 end