[discussion] CTMRG performances

[ogauthe]

Hello,

I recently started using PEPSKit. Kudos to all of you for all this work. I run finite temperature simulation and benchmark with my own python code (https://github.com/ogauthe/frostspin). I was happy to observe the results are nearly the same between frostspin and PEPSKit. However the performances were not as good as I was hoping for. I really think they can be improved, as I expect TensorKit to be much more efficient for the permutedims part.

Here I propose to give some benchmark points and find if some bottleneck can be improved in PEPSKit.

Setup: I use the simple update code to run a finite temperature simulation of the Heisenberg model on the square lattice. I add an ancilla leg to my PEPS tensor. As #184 is still a work in progress, I merge ancilla and physical leg together and define a Hamiltonian acting on both spaces to reuse PEPSKit code. I start from an infinite temperature product state, use simple update for imaginary time evolution. Once a given beta is reached, I use CTMRG to contract the tensor network and compute observables. I impose SU(2) symmetry. I can share the script if you want to reproduce.

Now there are of course several differences between the two codes. I think the most important ones are

• in frostspin I use an AB//BA unit cell, with only 2 sites.
• in the CTMRG, I cache enlarged corners to reuse them for the next iteration (PEPSKit simultaneous CTMRG looks a bit similar)
• I use a Lanczos based SVD and compute relatively few singular vectors
• when computing observables, I trace over the ancilla leg while contracting bra an ket tensors
• when computing observables, I compute reduce density matrices and normalize their trace (no norm computation)
• I avoid Zygote constrains on allocating memory (but I do not aggressively try to reuse it)
• due to different convergence criterion, PEPSKit does more iterations

Below are the logs for D=11, chi=121 to give precise values.
logs_profile_julia_1stnei_genx_CTM_SU2_D11_chi122_11b43a279c_c1.log
logs_profile_1stnei_genx_CTM_SU2_D11_chi121_rdm2nd_c7cb4eea37_c1.log

For D=11, chi=121, once compiling and unitary constructions are done

• 1 CTMRG step: frostpin takes ~9 sec for 4 sweeps (N-E-S-W) x 2 sites in the unit cell. PEPSKit takes ~42 sec, with a 4 site unit cell
• first neighbor bond energy: less than 1 sec on frostspin, 12 sec on PEPSKit
• second neighbor bond energy: 2 sec on frostspin, 25 sec on PEPSKit

The CTMRG iteration part looks reasonable, but finer profiling is needed. In frostspin, the bottleneck is always permute. Contractions are non negligible yet always cheaper. SVD are usually negligible. As I expect TensorKit permute to me much faster, there may be hidden bottleneck. On the other hand, observable computation is way slower than it has to be, although part of it is probably due to merged physical-ancilla legs.

A few comments:

• frostspin uses nearly 3 times less memory than PEPSKit (memory maximum is reached in 2nd neighbor density matrix computation).
• I get the same results with U(1) symmetry but frostspin abelian code has poor performances, benchmark are less relevant
• the same simple update code with Trivial sector gives incorrect results on PEPSKit. Still investigating.
• I plan to also benchmark the J1-J2 using Add 3-site simple update (aka 3-site cluster update) #171, but I have not been able to reproduce my previous results yet (this may be a PEPSKit bug, more probably just a script issue)
• I would be happy to provide more details on the algorithm or the benchmark
• My next step is to run a profiler on PEPSKit to have a better idea of time use.

[lkdvos]

Hi Olivier, as discussed, this is immensely helpful information, which at the very least means that it should be possible to (significantly) speed up several parts of our code.

At first glance, I would guess that this indicates some wrong contraction order somewhere. This is mostly since that would be the easiest thing to mess up, and we definitely did a bunch of tricky hackery to make that work (@autoopt is cool, but probably under-tested).

One thing to try could be to temporarily make a change here and insert a runtime check for the contraction order.

insert!(ex.args, length(ex.args), :(costcheck = :warn))

This effectively automatically inserts the costcheck=:warn in every tensor expression that is using @autoopt.
It might be interesting to see what comes out of that.
(Alternatively we could also simply check the contractions of the expectation_value functions)

[Yue-Zhengyuan]

@ogauthe (A bit unrelated) The finite temperature SU for Heisenberg model can be a great addition to PEPSKit docs. If you have time, could you please organize it into an example under examples/?

[ogauthe]

I have done some extensive testing of PEPSKit ctmrg_iterate. I find finite temperature to be a great benchmark as it is deterministic and easily reproducible.

Global bencharmks

I run the finite temperature code from infinite temperature (beta=0), time evolve it and compute CTMRG for β ∈ [1.5, 1.6, 1.7]. In all cases I use the setup described above with 2x2 unit cell with d=4 (merged physical + ancilla legs) and impose exactly 10 (17 for SU(2)) CTMRG iterations before computing observables. I saturate chi such that it stays very close to the target value D^2. Run were executed on a single thread on the same kind of node, same cluster. The nodes were shared with other jobs, but I expect the effects on performances to be small.

Trivial D=7 chi=49 ctm_niter = 10

PEPSKit Trivial D=7 chi=49: total time = 48 min, memory = 10GB
SimpleUpdate = 1 min (mostly compiling) ~ 2%
ctm_iteration = 39 min (compiling negligible) ~81%
observables = 8min (50 sec compiling, 7 min execution) ~17%
logs_profile_julia_1stnei_genx_CTM_Trivial_D7_chi49_ea0363a9ab_c1.log

frostspin Trivial D=7 chi=49: total time = 10 min, memory = 4GB
compiling time = negligible
SimpleUpdate = negligible
ctm_iteration = 500 sec (62% projector, 17% enlarged corner, 8% renormalize edge)
observables = 51 sec (6% 2nd neighbor, 4% 1st neighbor)
misc = 40 sec
logs_profile_1stnei_genx_CTM_trivial_D7_chi49_rdm2nd_c7cb4eea37_c1.log

U(1) D=11 chi=121 ctm_niter = 10

PEPSKit U(1) D=11 chi=121: total time = 7h, memory = 50GB
SimpleUpdate = 1 min (mostly compiling)
ctm_iteration = 6h (compiling negligible)
observables = 50min (~4 min compiling, 46 min execution)
logs_profile_julia_1stnei_genx_CTM_U1_D11_chi122_ea0363a9ab_c1.log

frostspin U(1) D=11 chi=121: total time = 93 min, memory = 15GB
compiling time ~ 2-3 min
SimpleUpdate = 1 min
ctm_iteration = 71 min (48% projector, 17% enlarged corner, 8% renormalize edge)
observables = 18 min (17% 2nd neighbor, 6% 1st neighbor)
logs_profile_1stnei_genx_CTM_U1_D11_chi121_rdm2nd_c7cb4eea37_c1.log

SU(2) D=11 chi=121 ctm_niter = 17

PEPSKit SU(2) D=11 chi=121: total time = 58 min, memory = 12GB
SimpleUpdate = 1 min (mostly compiling)
ctm_iteration = 36 min (compiling <1min)
observables = 19 min (4 min compiling, 46 min execution)
logs_profile_julia_1stnei_genx_CTM_SU2_D11_chi122_11b43a279c_c1.log

frostspin SU(2) D=11 chi=121: total time = 11min, memory = 4GB
compiling time ~ 2 min
SimpleUpdate = 1 min
ctm_iteration = 9 min (29% projector, 31% enlarged corner, 15% renormalize edge)
observables = <1 min (5% 2nd neighbor, 4% 1st neighbor)
logs_profile_1stnei_genx_CTM_SU2_D11_chi121_rdm2nd_87d18731ff_c1.log

I find PEPSKit is currently quite slower and uses much more memory, even taking into account a global factor 2 due to unit cell site number. This stays true for Trivial, U(1) and SU(2) symmetry. I use high level functions (simpleupdate and leading_boundary) with most arguments being defaults, so the comparison should be fair.
I believe most of the difference in memory use comes from the merged physical-ancilla legs in PEPSKit, while I can trace over the ancilla early in frostspin. This does not affect CTMRG iterations but should make observables more expensive on PEPSKit.

Local benchmark

On my workstation, I did some micro benchmark on the different steps of a CTMRG iterations. All the benchmarks are done on the same machine using IPython %timeit and julia RHEL @benchmark, again single thread. I first run the full code such that all compiling and unitary ("tree permuters") already took place, only run time is accounted. I checked the LRU cache is far from full. I benchmarked permutedims on a fully constructed enlarged corner with size chi^2*D^4; the largest intermediate tensor in contract_enlarged_corner has size a*d*chi^2*D^4.

SU(2) D=11 chi=121

PEPSKit SU(2) D=11 chi=121 single thread. GC is always negligible.
contract 1 enlarged corner: 890 ms
contract 1 half environment: 80 ms (=1 matmul)
construct 1 pair of projectors 1422 ms (includes 1 matmul & 1 full SVD)
renormalize 1 edge 860 ms
permute 1 enlarged corner 72 ms (size chi^2*D^4)

frostspin SU(2) D=11 chi=121 single thread
contract 1 enlarged corner: 320 ms
contract 1 half environment: 90 ms (=1 matmul)
construct 1 pair of projectors 430 ms (includes 3 matmul & 1 partial SVD)
renormalize 1 edge 330 ms
permute 1 enlarged corner 51 ms (size chi^2*D^4)

SU(2) D=16 chi=256

PEPSKit SU(2) D=16 chi=256 single thread. GC is always negligible.
contract 1 enlarged corner: 36 s
contract 1 half environment: 2.5 s (=1 matmul)
construct 1 pair of projectors 61 s (includes 1 matmul & 1 full SVD)
renormalize 1 edge 35.4 s
permute 1 enlarged corner 2.0 s (size chi^2*D^4)

frostspin SU(2) D=16 chi=256 single thread
contract 1 enlarged corner: 3.3 s
contract 1 half environment: 2.6 (=1 matmul)
construct 1 pair of projectors 11.8 s (includes 3 matmul & 1 partial SVD)
renormalize 1 edge 3.4 s
permute 1 enlarged corner 450 ms (size chi^2*D^4)

I am extremely surprised to find SU(2) permutedims to be faster with frostpin than with TensorKit. Once permutedims/reshape is done, contraction is nearly identical, which is not surprising as it boils down to (block) calls to gemm in both codes. Similarly a full SVD using gesdd takes nearly the same time, however frostspin default is Lanczos-based SVD which is much faster. I have not tried PEPSKit.IterSVD yet.
Note that even for D=16 chi-256, I did not reach the asymptotic behavior where the computing time is dominated by half environment / full environment matmul and the SVD, both scaling as O(chi^3*D^6)~O(D^12). Due to large prefactors and complexity of SU(2) implementation, the largest permutedims inside enlarge_corner is still a significant contribution despite scaling as O(a*d*chi^2*D^4)~O(D^8).

[Yue-Zhengyuan]

To confirm whether merging ancilla with physical legs is the main issue, do you think it's helpful to try a zero-temperature SU benchmark?

[Yue-Zhengyuan]

BTW, has the issue with 3-site simple update (for J1-J2) been resolved?

[ogauthe]

As shown in the benchmarks above, contracting an enlarged corner takes much longer than it should. It is implemented as @autoopt @tensor. It seems to me there is a bad interaction between these macros and symmetry implementation. When I write the contraction manually as a function I obtain the same result. For Trivial symmetry, @autoopt is slightly faster, however for U(1) it is already slower and for SU(2) its 4 times slower for D=16!

	TensorMap(Q,1)	contract_enlarged_corner(Q)	frostpsin
Trivial D=4 chi=16	2.0 ms	2.4 ms	1.2 ms
Trivial D=7 chi=49	288 ms	316 ms	212 ms
U(1) D=4 chi=16	1.5 ms	1.3 ms	2.5 ms
U(1) D=7 chi=49	54 ms	39 ms	55 ms
U(1) D=11 chi=121	2.4 s	1.8 s	2.3 s
SU(2) D=4 chi=16	8.8 ms	8.2 ms	2.6 ms
SU(2) D=7 chi=49	32 ms	20 ms	6.1 ms
SU(2) D=11 chi=121	933 ms	402 ms	229 ms
SU(2) D=16 chi=256	35 s	8 s	3.2 s

Full contract code

function contract_enlarged_corner(c1, t1, t4, (a, b))
    # SU2Irrep D=16 χ=256  total ~8 sec

    t4c1 = t4 * c1   # < 1 ms
    t1p = permute(t1, ((1,), (2, 3, 4)))  # < 1 ms

    #  C1---T1-6
    #  |    ||
    #  |    45
    #  |
    #  T4=2,3
    #  |
    #  1
    t4c1t1 = t4c1 * t1p   # 80 ms
    t4c1t1p = permute(t4c1t1, ((1, 3, 5, 6), (2, 4)))  # 1.0 s
    ap = permute(a, ((5, 2), (1, 3, 4)))  # < 1 ms
    t4c1t1a = t4c1t1p * ap  # 560 ms

    #  C1----T1-4
    #  |     ||
    #  |     |3
    #  |   5 |        1 2
    #  |    \|         \|
    #  T4----A-6      5-A*-3
    #  | \2  |          |
    #  1     7          4
    t4c1t1ap = permute(t4c1t1a, ((1, 4, 6, 7), (2, 3, 5)))  # 5.6 s
    bdagp = permute(b, ((3, 4), (5, 2, 1)))'  # < 1ms
    cnw = t4c1t1ap * bdagp  # 300 ms

    #  C1-T1-2 ---->4
    #  |  ||
    #  T4=AA*=3,5->5,6
    #  |  ||
    #  1  46
    #  1  23

    cnwp = permute(cnw, ((1, 4, 6), (2, 3, 5)))  # 680 ms
    return cnwp
end

function contract_enlarged_corner(Q::EnlargedCorner)
    return contract_enlarged_corner(Q.C, Q.E_2, Q.E_1, Q.A)
end

Q = PEPSKit.EnlargedCorner(network, env, (1, 1, 1))
ec1 = TensorMap(Q, 1)
ec2 = contract_enlarged_corner(Q)
@test ec1 ≈ ec2
@benchmark ec1 = TensorMap(Q, 1)
@benchmark ec2 = contract_enlarged_corner(Q)

I think this interaction macros/symmetries explains most of the timing difference above. For the global benchmark Trivial D=7 chi=49, once the factor 2 from the unit cell is added, CTMRG iterations take a similar time. frostspin is still a bit faster, probably due to Lanczos SVD vs full SVD. The difference in observables is probably due to the ancilla leg. As a side note, this issue may also explain why a separate benchmark on the spin-1 bilinear-biquadratic gives worse perf with SU(3) than with SU(2).

[ogauthe]

I created a dedicated repo to host the benchmark scripts: https://github.com/ogauthe/BenchmarkPEPS. You should be able to easily reproduce my results (both local and global).

@ogauthe (A bit unrelated) The finite temperature SU for Heisenberg model can be a great addition to PEPSKit docs. If you have time, could you please organize it into an example under examples/?

I should be able to adapt run_heisenberg.jl into an example inside PEPSKit.

To confirm whether merging ancilla with physical legs is the main issue, do you think it's helpful to try a zero-temperature SU benchmark?

A zero temperature SU benchmark makes sense, but is is much harder to produce. It is not really deterministic as it depends on the initial state and would be more sensitive to implementation details. Unfortunately translating a given quantum state from PEPSKit to frostspin is basically impossible due to different conventions. In the finite temperature case, we are saved thanks to a trivial, well defined initial state.

BTW, has the issue with 3-site simple update (for J1-J2) been resolved?

I haven’t revisited it yet.

[lkdvos]

Thanks for this detailed comparison, this shows exactly where to go and look for fixing this. It seems like the contraction order is indeed the culprit, and I know what to do next week. I'm still somewhat confused about the timings for the trivial cases, since that seems to indicate that at least for these dimensions the contraction order is correct, but it does make sense that the symmetries interact non-trivially with this.

I would start by trying to tweak the auto opt settings to see if that makes a difference, just filling in reasonable values for the dimensions and checking what comes out, but there might really be a bug in that code too, there clearly aren't any tests that check if the logic is actually happening as intended

[ogauthe]

Here are the virtual and corners spaces for SU(2):

• D=4: V_D = Rep[SU₂](0=>1, 1=>1); V_chi = Rep[SU₂](0=>2, 1=>3, 2=>1)
• D=7: V_D = Rep[SU₂](0=>1, 1=>2); V_chi = Rep[SU₂](0=>5, 1=>8, 2=>4)
• D=11: V_D = Rep[SU₂](0=>2, 1=>3); V_chi = Rep[SU₂](0=>12, 1=>19, 2=>9, 3=>1)
• D=16: V_D = Rep[SU₂](0=>2, 1=>3, 2=>1); V_chi = Rep[SU₂](0=>15, 1=>31, 2=>21, 3=>6)

For U(1) the SU(2) symmetry is more or less preserved

• D=4 V_D = Rep[U₁](0=>2, 1=>1, -1=>1); V_chi = Rep[U₁](0=>5, 1=>5, -1=>4, 2=>1, -2=>1)
• D=7: V_D = Rep[U₁](0=>3, 1=>2, -1=>2); V_chi = Rep[U₁](0=>17, 1=>11, -1=>11, 2=>5, -2=>5)
• D=11: V_D = Rep[U₁](0=>5, 1=>3, -1=>3); V_chi = Rep[U₁](0=>39, 1=>30, -1=>30, 2=>10, -2=>10, 3=>1, -3=>1)

[ogauthe]

I have also been running some experiments on the SU(3) 3-3 Heisenberg model at finite temperature. It can be implemented either directly as pure SU(3) or as the bilinear-biquadratic spin 1 Heisenberg model at the special point J = K. The spaces are then:

SU(2):

• V_d = Rep[SU₂](0=>1, 1=>1, 2=>1; S=1 and S=2 are degenerate and form SU(3) (2, 1, 0) while S=0 identifies as SU(3) singlet
• V_D = Rep[SU₂](0=>1, 1=>1, 2=>1; S=1 and S=2 are degenerate and form SU(3) (2, 1, 0)
• V_chi = Rep[SU₂](0=>3, 1=>6, 2=>7, 3=>2, 4=>1)

SU(3)

• V_d = Rep[SU₃]((1, 1, 0)=>1, (2, 0, 0)=>1)
• V_D = Rep[SU₃]((0, 0, 0)=>1, (2, 1, 0)=>1)
• V_chi = Rep[SU₃]((0, 0, 0)=>2, (2, 1, 0)=>4, (4, 2, 0)=>1)

The performances are pretty bad.

[Yue-Zhengyuan]

• I use a Lanczos based SVD and compute relatively few singular vectors

When there's additional symmetry, I think you need to know the dimension for each sector in the singular value spectrum in advance. May I ask if you determine them with a usual SVD first, and then switch to Lanczos SVD (IterSVD in PEPSKit)?