Skip to content

Add allocation-free pairwise contraction and cached contraction plan - #1762

Open
ipasichnyk wants to merge 2 commits into
ITensor:mainfrom
ipasichnyk:feature/alloc-free-contract
Open

Add allocation-free pairwise contraction and cached contraction plan#1762
ipasichnyk wants to merge 2 commits into
ITensor:mainfrom
ipasichnyk:feature/alloc-free-contract

Conversation

@ipasichnyk

Copy link
Copy Markdown

Description

Adds an allocation-free path for repeated pairwise tensor contractions. The standard A * B / contract!! path allocates fresh permute and output buffers on every call. In hot loops that contract the same network shape many times (e.g. belief-propagation / iterative simulation), and especially on GPUs, these per-call device allocations dominate runtime.
This PR adds:

  • NDTensors.ContractScratch plus _contract_prealloc! / contract_prealloc!, which contract into a pre-allocated output using cached permute scratch (Ap/Bp/Cp) that is lazily sized on first use and reused afterward.
  • An ITensor-level contract_prealloc!(scr, C, A, B) that writes into C's existing storage.
  • ContractionPlan / contraction_plan(sequence) which mirrors a binary contraction sequence: the first run allocates the intermediate buffers, and later runs replay allocation-free via contract!(C, As, plan).
    The change is purely additive; no existing public interfaces are modified. The permuteC path is supported only for β == 0 (it errors otherwise), which is the only case the plan needs. No new dependencies are required.

Fixes #(issue)

Minimal demonstration of previous behavior

julia
using ITensors
i, j, k, l = Index(3, "i"), Index(4, "j"), Index(5, "k"), Index(2, "l")
A = random_itensor(i, j, k)
B = random_itensor(k, l)
# Every contraction allocates fresh permute/output buffers, on every iteration:
for _ in 1:3
    C = A * B
    @show @allocated A * B   # > 0 on each call
end

Minimal demonstration of new behavior

using ITensors.NDTensors: ContractScratch
i, j, k, l = Index(3, "i"), Index(4, "j"), Index(5, "k"), Index(2, "l")
A = random_itensor(i, j, k)
B = random_itensor(k, l)
C   = ITensor(0.0, (i, j, l))   # pre-allocated output
scr = ContractScratch()         # reused permute scratch
ITensors.contract_prealloc!(scr, C, A, B)   # first call sizes the buffers
@assert C ≈ A * B
# Subsequent calls reuse the buffers, no new allocations for the contraction:
for _ in 1:3
    @show @allocated ITensors.contract_prealloc!(scr, C, A, B)   # ~0
end
# Cached sequenced plan: first run allocates intermediates, later runs replay allocation-free.
D    = random_itensor(l, j)
As   = ITensor[A, B, D]
out  = ITensor(0.0, inds((A * B) * D))
plan = ITensors.contraction_plan([[1, 2], 3])
ITensors.contract!(out, As, plan)   # warms up (allocates intermediates)
ITensors.contract!(out, As, plan)   # allocation-free replay
@assert out ≈ (A * B) * D

How Has This Been Tested?

Added a new @testset "Allocation-free contraction ($T)" (for T in (Float64, ComplexF64)) to test/base/test_contract.jl, which is auto-included by test/base/runtests.jl. All base contraction tests pass along with the new ones.

contract_prealloc! produces the same result as A * B, and repeated calls reusing the same ContractScratch stay correct.
contract_prealloc! is correct when the output index order forces the permuteC path.
ContractionPlan matches sequential (A * B) * D, and a second replay into the same reused output buffer is still correct.

Checklist:

  • My code follows the style guidelines of this project. Please run the ITensorFormatter in the base directory of the repository (~/.julia/dev/ITensors) to format your code according to our style guidelines.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that verify the behavior of the changes I made.
  • I have made corresponding changes to the documentation.
  • My changes generate no new warnings.
  • Any dependent changes have been merged and published in downstream modules.

ipasichnyk and others added 2 commits July 15, 2026 15:21
NDTensors to contract into pre-allocated output and permute buffers, avoiding
per-call device allocations. Add ITensor-level contract_prealloc! and a
ContractionPlan that mirrors a binary contraction sequence, allocating
intermediates on first run and reusing them for allocation-free replay.
Co-authored-by: Cursor <cursoragent@cursor.com>
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