Add allocation-free pairwise contraction and cached contraction plan - #1762
Open
ipasichnyk wants to merge 2 commits into
Open
Add allocation-free pairwise contraction and cached contraction plan#1762ipasichnyk wants to merge 2 commits into
ipasichnyk wants to merge 2 commits into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.ContractScratchplus_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.contract_prealloc!(scr, C, A, B)that writes intoC'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 viacontract!(C, As, plan).The change is purely additive; no existing public interfaces are modified. The
permuteCpath 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
Minimal demonstration of new behavior
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:
~/.julia/dev/ITensors) to format your code according to our style guidelines.