[WIP] Allow CTMRG methods to use an allocator

Project.toml

@@ -33,6 +33,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
+[sources]
+TensorKit = {url = "https://github.com/QuantumKitHub/TensorKit.jl", rev = "main"}
+
 [extensions]
 PEPSKitAdaptExt = "Adapt"
 PEPSKitGPUArraysExt = "GPUArrays"

src/PEPSKit.jl

@@ -52,7 +52,7 @@ import MPSKitModels:
     transverse_field_ising, heisenberg_XYZ, heisenberg_XXZ,
     hubbard_model, bose_hubbard_model, tj_model
 using FiniteDifferences
-using OhMyThreads: tmap, tmap!, SerialScheduler
+using OhMyThreads: tmap, tmapreduce, tmap!, SerialScheduler
 using DocStringExtensions
 
 include("Defaults.jl")  # Include first to allow for docstring interpolation with Defaults values

src/algorithms/contractions/ctmrg/gaugefix.jl

@@ -17,9 +17,9 @@ Multiply corner tensor with incoming and outgoing gauge signs.
 ```
 """
 function fix_gauge_corner(
-        corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor
+        corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor, allocator
     )
-    return @autoopt @tensor corner_fix[χ_out; χ_in] :=
+    return @autoopt @tensor allocator = allocator corner_fix[χ_out; χ_in] :=
         σ_out[χ_out; χ1] * corner[χ1; χ2] * conj(σ_in[χ_in; χ2])
 end
 
@@ -28,11 +28,12 @@ $(SIGNATURES)
 
 Apply `fix_gauge_corner` to the northwest corner with appropriate row and column indices.
 """
-function fix_gauge_northwest_corner((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_northwest_corner((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_corner(
         corner(env, NORTHWEST, row, col),
         signs[WEST, row, col],
         signs[NORTH, row, _next(col, end)],
+        allocator,
     )
 end
 
@@ -41,11 +42,12 @@ $(SIGNATURES)
 
 Apply `fix_gauge_corner` to the northeast corner with appropriate row and column indices.
 """
-function fix_gauge_northeast_corner((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_northeast_corner((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_corner(
         corner(env, NORTHEAST, row, col),
         signs[NORTH, row, col],
         signs[EAST, _next(row, end), col],
+        allocator,
     )
 end
 
@@ -54,11 +56,12 @@ $(SIGNATURES)
 
 Apply `fix_gauge_corner` to the southeast corner with appropriate row and column indices.
 """
-function fix_gauge_southeast_corner((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_southeast_corner((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_corner(
         corner(env, SOUTHEAST, row, col),
         signs[EAST, row, col],
         signs[SOUTH, row, _prev(col, end)],
+        allocator,
     )
 end
 
@@ -67,11 +70,12 @@ $(SIGNATURES)
 
 Apply `fix_gauge_corner` to the southwest corner with appropriate row and column indices.
 """
-function fix_gauge_southwest_corner((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_southwest_corner((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_corner(
         corner(env, SOUTHWEST, row, col),
         signs[SOUTH, row, col],
         signs[WEST, _prev(row, end), col],
+        allocator,
     )
 end
 
@@ -88,7 +92,7 @@ Multiply edge tensor with incoming and outgoing gauge signs.
 ```
 """
 @generated function fix_gauge_edge(
-        edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor
+        edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor, allocator
     ) where {T, S, N}
     edge_fix_e = tensorexpr(
         :edge_fix,
@@ -102,7 +106,7 @@ Multiply edge tensor with incoming and outgoing gauge signs.
     σ_in_e = tensorexpr(:σ_in, (envlabel(:in),), (envlabel(2),))
     return macroexpand(
         @__MODULE__,
-        :(return @autoopt @tensor $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)),
+        :(return @autoopt @tensor allocator = allocator $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)),
     )
 end
 
@@ -111,9 +115,9 @@ $(SIGNATURES)
 
 Apply `fix_gauge_edge` to the north edge with appropriate row and column indices.
 """
-function fix_gauge_north_edge((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_north_edge((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_edge(
-        edge(env, NORTH, row, col), signs[NORTH, row, col], signs[NORTH, row, _next(col, end)],
+        edge(env, NORTH, row, col), signs[NORTH, row, col], signs[NORTH, row, _next(col, end)], allocator
     )
 end
 
@@ -122,9 +126,9 @@ $(SIGNATURES)
 
 Apply `fix_gauge_edge` to the east edge with appropriate row and column indices.
 """
-function fix_gauge_east_edge((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_east_edge((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_edge(
-        edge(env, EAST, row, col), signs[EAST, row, col], signs[EAST, _next(row, end), col]
+        edge(env, EAST, row, col), signs[EAST, row, col], signs[EAST, _next(row, end), col], allocator
     )
 end
 
@@ -133,9 +137,9 @@ $(SIGNATURES)
 
 Apply `fix_gauge_edge` to the south edge with appropriate row and column indices.
 """
-function fix_gauge_south_edge((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_south_edge((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_edge(
-        edge(env, SOUTH, row, col), signs[SOUTH, row, col], signs[SOUTH, row, _prev(col, end)],
+        edge(env, SOUTH, row, col), signs[SOUTH, row, col], signs[SOUTH, row, _prev(col, end)], allocator
     )
 end
 
@@ -144,9 +148,9 @@ $(SIGNATURES)
 
 Apply `fix_gauge_edge` to the west edge with appropriate row and column indices.
 """
-function fix_gauge_west_edge((row, col), env::AbstractCTMRGEnv, signs)
+function fix_gauge_west_edge((row, col), env::AbstractCTMRGEnv, signs, allocator)
     return fix_gauge_edge(
-        edge(env, WEST, row, col), signs[WEST, row, col], signs[WEST, _prev(row, end), col]
+        edge(env, WEST, row, col), signs[WEST, row, col], signs[WEST, _prev(row, end), col], allocator
     )
 end
 

src/algorithms/contractions/ctmrg/network_value.jl

@@ -10,9 +10,9 @@ function _contract_site(
         C_northwest, C_northeast, C_southeast, C_southwest,
         E_north::CTMRG_PEPS_EdgeTensor, E_east::CTMRG_PEPS_EdgeTensor,
         E_south::CTMRG_PEPS_EdgeTensor, E_west::CTMRG_PEPS_EdgeTensor,
-        O::PEPSSandwich,
+        O::PEPSSandwich, allocator
     )
-    return @autoopt @tensor E_west[χ_WSW D_W_above D_W_below; χ_WNW] *
+    return @autoopt @tensor allocator = allocator E_west[χ_WSW D_W_above D_W_below; χ_WNW] *
         C_northwest[χ_WNW; χ_NNW] *
         E_north[χ_NNW D_N_above D_N_below; χ_NNE] *
         C_northeast[χ_NNE; χ_ENE] *
@@ -27,9 +27,9 @@ function _contract_site(
         C_northwest, C_northeast, C_southeast, C_southwest,
         E_north::CTMRG_PF_EdgeTensor, E_east::CTMRG_PF_EdgeTensor,
         E_south::CTMRG_PF_EdgeTensor, E_west::CTMRG_PF_EdgeTensor,
-        O::PFTensor,
+        O::PFTensor, allocator
     )
-    return @autoopt @tensor E_west[χ_WSW D_W; χ_WNW] *
+    return @autoopt @tensor allocator = allocator E_west[χ_WSW D_W; χ_WNW] *
         C_northwest[χ_WNW; χ_NNW] *
         E_north[χ_NNW D_N; χ_NNE] *
         C_northeast[χ_NNE; χ_ENE] *
@@ -43,7 +43,7 @@ end
 @generated function _contract_site(
         C_northwest, C_northeast, C_southeast, C_southwest,
         E_north::TE, E_east::TE, E_south::TE, E_west::TE,
-        O::PEPOSandwich{H},
+        O::PEPOSandwich{H}, allocator
     ) where {TE <: CTMRGEdgeTensor, H}
     @assert numout(TE) == H + 3
 
@@ -67,15 +67,15 @@ end
         pepo_es...,
     )
 
-    return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs))
+    return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs))
 end
 
 """
-    _contract_site(ind::Tuple{Int,Int}, network::InfiniteSquareNetwork, env::AbstractCTMRGEnv)
+    _contract_site(ind::Tuple{Int,Int}, network::InfiniteSquareNetwork, env::AbstractCTMRGEnv, allocator)
 
 Contract around a single site `ind` of a square network using a given CTMRG environment.
 """
-function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::CTMRGEnv)
+function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::CTMRGEnv, allocator)
     r, c = ind
     return _contract_site(
         corner(env, NORTHWEST, r - 1, c - 1),
@@ -84,11 +84,11 @@ function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, en
         corner(env, SOUTHWEST, r + 1, c - 1),
         edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1),
         edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1),
-        network[r, c],
+        network[r, c], allocator
     )
 end
 
-function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::C4vCTMRGEnv)
+function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::C4vCTMRGEnv, allocator)
     r, c = ind
     @assert r == c == 1 lazy"only trivial unit cells are supported for C4v-CTMRG, ind ($ind) must be (1, 1)"
     return _contract_site(
@@ -98,32 +98,34 @@ function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, en
         corner(env, SOUTHWEST, r + 1, c - 1),
         edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1),
         edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1),
-        network[r, c],
+        network[r, c], allocator,
     )
 end
 
 ## Normalization contractions
 function _contract_corners(
         C_northwest::CTMRGCornerTensor, C_northeast::CTMRGCornerTensor,
         C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor,
+        allocator,
     )
-    return @tensor C_northwest[1; 2] * C_northeast[2; 3] *
+    return @tensor allocator = allocator C_northwest[1; 2] * C_northeast[2; 3] *
         C_southeast[3; 4] * C_southwest[4; 1]
 end
 
 """
-    _contract_corners(ind::Tuple{Int,Int}, env::CTMRGEnv)
+    _contract_corners(ind::Tuple{Int,Int}, env::CTMRGEnv, allocator)
 
 Contract all corners around the south-east at position `ind` of the CTMRG
 environment `env`.
 """
-function _contract_corners(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv)
+function _contract_corners(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator)
     r, c = ind
     return _contract_corners(
         corner(env, NORTHWEST, r - 1, c - 1),
         corner(env, NORTHEAST, r - 1, c),
         corner(env, SOUTHEAST, r, c),
         corner(env, SOUTHWEST, r, c - 1),
+        allocator,
     )
 end
 
@@ -132,6 +134,7 @@ end
         C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor,
         E_east::CTMRGEdgeTensor{T, S, N},
         E_west::CTMRGEdgeTensor{T, S, N},
+        allocator,
     ) where {T, S, N}
     C_northwest_e = tensorexpr(:C_northwest, (envlabel(:NW),), (envlabel(:N),))
     C_northeast_e = tensorexpr(:C_northeast, (envlabel(:N),), (envlabel(:NE),))
@@ -150,7 +153,7 @@ end
         E_west_e, C_northwest_e, C_northeast_e, E_east_e, C_southeast_e, C_southwest_e,
     )
 
-    return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs))
+    return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs))
 end
 
 """
@@ -159,7 +162,7 @@ end
 Contract the vertical edges and corners around the east edge at position `ind` of the
 CTMRG environment `env`.
 """
-function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv)
+function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator)
     r, c = ind
     return _contract_vertical_edges(
         corner(env, NORTHWEST, r - 1, c - 1),
@@ -168,13 +171,15 @@ function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv)
         corner(env, SOUTHWEST, r + 1, c - 1),
         edge(env, EAST, r, c),
         edge(env, WEST, r, c - 1),
+        allocator
     )
 end
 
 @generated function _contract_horizontal_edges(
         C_northwest::CTMRGCornerTensor, C_northeast::CTMRGCornerTensor,
         C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor,
         E_north::CTMRGEdgeTensor{T, S, N}, E_south::CTMRGEdgeTensor{T, S, N},
+        allocator,
     ) where {T, S, N}
     C_northwest_e = tensorexpr(:C_northwest, (envlabel(:W),), (envlabel(:NW),))
     C_northeast_e = tensorexpr(:C_northeast, (envlabel(:NE),), (envlabel(:E),))
@@ -193,16 +198,16 @@ end
         C_northwest_e, E_north_e, C_northeast_e, C_southeast_e, E_south_e, C_southwest_e,
     )
 
-    return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs))
+    return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs))
 end
 
 """
-    _contract_horizontal_edges(ind::Tuple{Int,Int}, env::CTMRGEnv)
+    _contract_horizontal_edges(ind::Tuple{Int,Int}, env::CTMRGEnv, allocator)
 
 Contract the horizontal edges and corners around the south edge at position `ind` of the
 CTMRG environment `env`.
 """
-function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv)
+function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator)
     r, c = ind
     return _contract_horizontal_edges(
         corner(env, NORTHWEST, r - 1, c - 1),
@@ -211,5 +216,6 @@ function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv)
         corner(env, SOUTHWEST, r, c - 1),
         edge(env, NORTH, r - 1, c),
         edge(env, SOUTH, r, c),
+        allocator,
     )
 end