Repository navigation
Refactor simple update and reduce weight-absorption allocations - #459
Conversation
|
Some of the refactors are peeled off from #144, so this one should go first. |
Codecov Report❌ Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
Rework the nearest-neighbor simple update along the lines of the Canopy.jl two-site kernel: - Absorb diagonal bond weights directly on each fusion-tree block, in the native leg order, with a single fused Strided broadcast. This replaces the permute-multiply-permute loop and the `input_axes`/`output_perm` keyword arguments of `_absorb_weight`. - Update a bond without rotating the tensors: derive the legs from the bond direction, split each site with one permutation and a QR (`_bond_qr`), apply the gate to the reduced tensors (`_apply_bond_gate`), and contract the result straight back into the native leg order (`_undo_bond_qr`). Orient the bond such that the SVD reproduces the original arrow, so no flips are needed. - Serve temporaries from `MPSKit.default_allocator`, created once per `su_iter`. Temporaries never outlive the function that allocates them. Remove the now unused `_bond_rotation`, `_absorb_weight!`, `_permute_cluster`, `_permute_to_last` and `_apply_gate`, add `absorb_weight!` and `absorb_s!`, and require TensorOperations 5.8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add a `timeevol` benchmark module that times one `su_iter` sweep over a 2 x 2 unit cell for trivial and U(1) Heisenberg PEPS, fermionic Hubbard PEPS and a Heisenberg PEPO, at two bond dimensions each. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Per fusion-tree block, `block(wt, c)` and `t[f₁, f₂]` each redo a lookup in TensorKit's global structure cache, which dominated the cost of absorbing weights into tensors with many small blocks. Look up the weight blocks once per call and iterate with `subblocks` instead. Also make the tensor rank static in the `ntuple` that reshapes the weights, which was otherwise not inferred and caused dynamic dispatch on every broadcast. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
a8ae4e3 to
0e1c18e
Compare
|
I pushed a few commits on top of this PR, taking inspiration from the two-site kernel in Canopy.jl, and rebased the branch onto Changes
Benchmarks One
Allocation counts ( Testing I checked that the new kernel reproduces the 🤖 Generated with Claude Code |
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
I tried to include some cleanup as well, but might have removed some of the helper functions you still needed, I do apologize if this is the case, let me know which and I'll try adding them back. Otherwise, I'd say this is definitely an improvement, and overall the code didn't really get all that much more complicated, the added lines are mostly due to adding the benchmark cases so I'd say that's not too bad.
Otherwise good to go from my end!
|
I'll see if I can adapt #144 cleanly before deciding if any of the rotation helpers need to stay. |
This PR further improves performance of simple update, with a few more refactoring.
_absorb_weightaccept more flexible tensor leg order, avoiding intermediate permutations between PEPS/PEPO and MPS (cluster) order.iter/tkeywords to restarttimestepthrough the public API.