Repository navigation
Conversation
`makeblascontractable` received the scalar type of the destination as an unconstrained argument, which Julia does not specialize on. As a result, the allocation and permutation of inputs that are not directly BLAS-contractable were dynamically dispatched, which also affects the `@tensor` macro. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The index notation files are included in a module-level `@nospecialize` block to reduce the compilation time of the macros. However, `ncontree.jl` and `indexordertree.jl` contain functions that are also called at runtime by `ncon`, such as `check_nconstyle`, `ncontree` and `indexordertree`, which therefore ran without specialization on their arguments. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- `check_nconstyle` and `_verify_indexorder` count label occurrences instead of repeatedly searching and deleting from an abstractly typed vector. - `ncontree` and `indexordertree` share a single implementation that merges partial trees in a type-stable manner, and represent the contraction tree as a list of pairwise contractions. The nested trees that are returned are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The network and the intermediate results of `ncon` have non-uniform types, which previously
caused type instabilities and runtime type inference throughout the evaluation of the
contraction tree. Now:
- `resolve_traces` canonicalizes the tensors and labels to a `Vector{Any}` and
`Vector{Vector{Int}}`, and the output and conjugation flags are canonicalized as well,
such that the remainder of `ncon` is type stable and compiled only once.
- The contraction tree is traversed as a list of pairwise contractions with integer ids.
- The keyword arguments are processed once, instead of for every pairwise contraction.
- All label manipulations happen on `Vector{Int}`, and every pairwise contraction goes
through a single function barrier, in which the scalar type is inferred statically
instead of through `promote_op` at runtime.
As a side effect, `ncon` can now be differentiated with Zygote, and networks with traces
can be specified with row matrices as index lists.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The tuple implementation of `contract_indices` manipulates the labels in a type-unstable manner, and is compiled for every combination of numbers of indices. For labels stored in vectors, as in `ncon`, the permutations are now computed as vectors in a type-stable manner, and only converted to tuples at the end. `ncon` passes these tuples to its function barrier as separate arguments rather than as `Index2Tuple`s, which avoids boxing them. This reduces the compilation time for networks with different numbers of indices, at a small cost in runtime for small networks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`makeblascontractable` returned either its input or a permuted copy, whose types differ when
the copy is allocated by an allocator such as Bumper's. The resulting union is not an
`isbits` type and was therefore heap allocated on every contraction, even when no copy was
required. Its specialization on the scalar type also made Enzyme fail ("could not determine
how to keep a pointer phi rooted").
`_blas_contract!` now makes the inputs BLAS-contractable one at a time, passing them on to
the next step in each branch rather than returning them, such that the union is never
materialized. Each input is reshaped into a matrix before it is passed on, so that the
subsequent steps no longer depend on the number of indices of the inputs, which keeps the
number of compiled methods (and the compilation time) at the level of before this change.
`makeblascontractable` and `_unsafe_blas_contract!` are no longer used and are removed,
together with the tests of the former.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
... and 5 files with indirect coverage changes 🚀 New features to boost your workflow:
|
This branch has not been deployed
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
1 participant
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.
nconspent most of its time on bookkeeping rather than on contractions: runtime type inference (promote_op) at every node of the contraction tree, abstractly typed containers, repeated keyword argument handling, and unspecialized network analysis. This PR makes the label handling type stable, so that every pairwise contraction goes through a single dynamic dispatch. It also includes the fix for #273, previously #315.makeblascontractablereceived the scalar type as an unconstrained argument, which Julia does not specialize on, so the returned tensor was inferred asAnyand the rest of the BLAS contraction path was dynamically dispatched. It also returned either its input or a permuted copy, whose types differ when the copy comes from an allocator such as Bumper's, and that union was heap allocated on every contraction.makeblascontractableis removed:_blas_contract!now makes the inputs BLAS-contractable one at a time and passes them on within each branch, so the union is never materialized. Each input is reshaped into a matrix before it is passed on, so the later steps do not depend on its number of indices, which keeps the number of compiled methods down. This also speeds up@tensor(see the reference row below).ncontree.jlandindexordertree.jlwere included in the module-level@nospecializeblock meant for the macros, but are also used at runtime byncon.ncontreeandindexordertreeare unchanged.ncon. Tensors, labels, output and conjugation flags are canonicalized once, all label manipulations happen onVector{Int}, and keyword arguments are processed once.contract_indicesfor labels stored in vectors. This avoids the type-unstable tuple implementation, which is compiled for every combination of numbers of indices. It also applies totensorcontract(IC, A, IA, B, IB)with vector labels.As a side effect,
nconcan now be differentiated with Zygote (including networks with traces), and traces can be specified with row matrices as index lists.Fixes #273, together with QuantumKitHub/Strided.jl#94. Allocations in its reproducer (
StridedBLAS, Bumper allocator, latest Strided and StridedViewsmain):mainFloat64ComplexF64ComplexF64,conjBenchmarks
Exclusive node, Julia 1.13.1, a single thread, every measurement in a fresh process, registered Strided. Tensors have dimension 2 per index unless noted, so the overhead dominates. Minimum times; three interleaved rounds agree to within 2%.
mainTupleinputorderkeywordComplexF64@tensorreference (first row)nconcall / with new typesThe last row is from an earlier run with the same setup (the previous version of this PR, without the #273 fix); the final version matches it within noise in local comparisons (median 4.32 s without and 4.39 s with the fix).
🤖 Generated with Claude Code