Skip to content

Cache TagSet("Site") in v3's apply_mpo() to speed up KPM and other MPO-heavy paths - #95

Merged
joselado merged 2 commits into
masterfrom
worktree-kpm-v3-tagset-perf
Jul 30, 2026
Merged

Cache TagSet("Site") in v3's apply_mpo() to speed up KPM and other MPO-heavy paths#95
joselado merged 2 commits into
masterfrom
worktree-kpm-v3-tagset-perf

Conversation

@joselado

@joselado joselado commented Jul 29, 2026

Copy link
Copy Markdown
Owner

Summary

  • mpscpp3/chain_session.h's apply_mpo() (both overloads) re-parsed TagSet("Site") from scratch on every call via out.noPrime(TagSet("Site")). apply_mpo() is the single most-invoked MPO-application helper in the file: every Chebyshev moment of KPM's dynamical-correlator recursion (kpm_moments_full/kpm_moments_accelerated) calls it once or twice, and it also backs time evolution, the generalized-eigenproblem solver, and CVM's bicstab solver.
  • The exact same class of overhead (TagSet construction re-parsing its string argument every call) was already profiled with valgrind --tool=callgrind and fixed elsewhere in this file (arnoldi_smallest_real, idmrg_extend_HL/idmrg_extend_HR, PR Wire ITensor v3 C++ iDMRG backend into infinitechain.py #87's follow-up) via a function-local static const TagSet — but that pass never touched apply_mpo() itself. This applies the identical, already-validated pattern there. Zero semantic change (same parsed tag value either way).
  • Measured directly on the actual KPM hot loop: ~2-4% faster when bond dimension/moment cost is small enough that applyMPO's own SVD work doesn't dominate the per-call fixed overhead; negligible when it does (unlike the iDMRG Arnoldi case's 6-8%, since that loop spends proportionally more time in cheap per-step ITensor operations relative to the parsing cost). v2 is unaffected — it predates the Tags system.
  • Also fixed examples/dynamical_correlator/dynamical_correlator_v2_VS_v3/main.py, found broken while benchmarking this change: it used a stale i=i,j=j,name="ZZ" shorthand that get_dynamical_correlator no longer accepts (name= must be a pair of MultiOperators), so it raised RuntimeError: No active exception to reraise on both backends. Replaced with name=(sc.Sz[i],sc.Sz[j]), the same convention tests/test_dynamical_correlator.py already uses.

Test plan

  • pytest tests/test_dynamical_correlator.py -k "kpm or cvm" (parametrized over itensor_version incl. 3) passes before and after the change
  • Full suite: pytest tests — 268 passed, 9 skipped, unchanged from baseline
  • Benchmarked before/after with a standalone script exercising get_dynamical_correlator(mode="DMRG", submode="KPM") at multiple chain sizes/moment counts, built via the fast worktree-local mpscpp3 rebuild trick (reuse libitensor.a, just make pybind)
  • Fixed example re-run end-to-end against both real compiled backends (mpscpp2 + mpscpp3, not the ED fallback), stable timings across repeated runs

🤖 Generated with Claude Code

https://claude.ai/code/session_01DoZKWzhMo7qKoiP12gbcvv

joselado and others added 2 commits July 29, 2026 18:42
…vy paths

apply_mpo() re-parsed TagSet("Site") from scratch on every call via
out.noPrime(TagSet("Site")) -- the same class of overhead already found
and fixed (via a function-local static TagSet) in arnoldi_smallest_real
and idmrg_extend_HL/HR, but never applied to this helper. apply_mpo() is
the single most-invoked MPO-application wrapper in chain_session.h: every
Chebyshev moment of KPM's dynamical-correlator recursion calls it once or
twice, and it also backs time evolution, the generalized-eigenproblem
solver, and CVM's bicstab solver. Zero semantic change (same parsed tag
value either way).

Measured on the actual KPM hot loop (many small-bond-dimension moments,
where applyMPO's own SVD cost doesn't dominate): ~2-4% faster. Negligible
when bond dimension/moment cost is large enough that applyMPO's own work
dominates the per-call fixed overhead -- unlike the iDMRG Arnoldi case
(6-8%), which spends much more of its time in cheap per-step ITensor
operations relative to the parsing cost.

tests/test_dynamical_correlator.py's kpm/cvm itensor_version=3 tests and
the full suite (268 passed, 9 skipped) pass unchanged after the change.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DoZKWzhMo7qKoiP12gbcvv
…ample

get_dynamical_correlator's name= parameter has required a pair of
MultiOperators for a while now (see kpmdmrg.py's own type check on
name[0]/name[1]); this example still passed the old i=i,j=j,name="ZZ"
string shorthand, so running it raised "RuntimeError: No active
exception to reraise" on both backends. Replaced with name=(sc.Sz[i],
sc.Sz[j]), the same convention tests/test_dynamical_correlator.py
already uses. Verified it now runs end-to-end against both real compiled
backends (mpscpp2 and mpscpp3, not the ED fallback) and produces stable
timings across repeated runs.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DoZKWzhMo7qKoiP12gbcvv
@joselado
joselado marked this pull request as ready for review July 30, 2026 14:32
@joselado
joselado merged commit 15aa9de into master Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant