Reduce the overhead of ncon

docs/src/changelog.md

@@ -17,8 +17,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `ncon` can be differentiated with Zygote, including networks with traces.
+- `ncon` accepts traces in networks whose index lists are given as row matrices, e.g. `[1 1 -1]`.
+- Contractions using the `StridedBLAS` backend no longer heap allocate a view of an input tensor when an allocator such as Bumper's is used ([#273](https://github.com/QuantumKitHub/TensorOperations.jl/issues/273)).
+
 ### Performance
 
+- `ncon` has significantly less overhead: the network analysis and contraction tree construction are type stable, and every pairwise contraction goes through a single dynamic dispatch.
+- Contractions of `Array`s that require a permuted copy of an input no longer suffer from a type instability.
+- Index labels given as vectors, as in `ncon` or the method syntax `tensorcontract(IC, A, IA, B, IB)`, are converted to index permutations in a type-stable manner, which avoids compiling the conversion for every combination of numbers of indices.
 - `ncon` partitions its intermediates as `(open indices of A; open indices of B)` instead of placing all indices in the codomain, which avoids repartitioning copies for tensor types such as `TensorMap`s with symmetries ([#311](https://github.com/QuantumKitHub/TensorOperations.jl/issues/311)).
 
 ## [5.8.2](https://github.com/QuantumKitHub/TensorOperations.jl/compare/v5.8.1...v5.8.2) - 2026-09-30

ext/TensorOperationsChainRulesCoreExt.jl

@@ -23,6 +23,7 @@ using LinearAlgebra
 @non_differentiable TensorOperations.tensoralloc_contract(args...)
 @non_differentiable TensorOperations.promote_contract(args...)
 @non_differentiable TensorOperations.promote_add(args...)
+@non_differentiable TensorOperations.contract_indices(args...)
 
 # Cannot free intermediate tensors when using AD
 # Thus we change the forward passes: `istemp=false` and `tensorfree!` is a no-op
@@ -193,5 +194,12 @@ end
 @non_differentiable TensorOperations.nconoutput(args...)
 @non_differentiable TensorOperations.check_nconstyle(args...)
 @non_differentiable TensorOperations.indexordertree(args...)
+@non_differentiable TensorOperations._ncontree_steps(args...)
+@non_differentiable TensorOperations._indexordertree_steps(args...)
+# label manipulations that mutate arrays, which is not supported by e.g. Zygote
+@non_differentiable TensorOperations.unique2(args...)
+@non_differentiable TensorOperations._ncon_labelvector(args...)
+@non_differentiable TensorOperations._ncon_conjlist(args...)
+@non_differentiable TensorOperations._ncon_symdiff(args...)
 
 end

src/TensorOperations.jl

@@ -51,8 +51,11 @@ include("indexnotation/analyzers.jl")
 include("indexnotation/preprocessors.jl")
 include("indexnotation/postprocessors.jl")
 include("indexnotation/contractiontrees.jl")
+# also used at runtime by `ncon`, so these should be specialized
+@specialize
 include("indexnotation/ncontree.jl")
 include("indexnotation/indexordertree.jl")
+@nospecialize
 include("indexnotation/poly.jl")
 include("indexnotation/optdata.jl")
 include("indexnotation/optimaltree.jl")

src/implementation/blascontract.jl

@@ -41,74 +41,57 @@ function blas_contract!(
     return C
 end
 
-# implement necessary permutations
+# implement necessary permutations - each input is made BLAS-contractable and reshaped into a
+# matrix in a separate step within each branch, to avoid boxing/type instabilities from
+# returning a union of the input and its permuted copy; as the subsequent steps only receive
+# matrices, they no longer depend on the number of indices of the inputs
 function _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
-    TC = eltype(C)
-
-    A_, pA, flagA = makeblascontractable(A, pA, TC, backend, allocator)
-    B_, pB, flagB = makeblascontractable(B, pB, TC, backend, allocator)
-
-    # the partition of `ipAB` is required by `isblasdestination` and `tensoralloc_add`, but
-    # the actual contraction only needs the linearized permutation
     ipAB = oindABinC(pAB, pA, pB)
-    ipAB′ = linearize(ipAB)
-    flagC = isblasdestination(C, ipAB)
-    if flagC
-        C_ = C
-        _unsafe_blas_contract!(C_, A_, pA, B_, pB, ipAB′, α, β)
+    if isblascontractable(A, pA) && eltype(A) == eltype(C)
+        return _blas_contract_B!(C, _blas_matrix(A, pA), B, pB, pAB, ipAB, α, β, backend, allocator)
+    else
+        A_ = tensoralloc_add(eltype(C), A, pA, false, Val(true), allocator)
+        Anew = tensoradd!(SV(A_), A, pA, false, One(), Zero(), backend, allocator)
+        Amat = _blas_matrix(Anew, trivialpermutation(pA))
+        _blas_contract_B!(C, Amat, B, pB, pAB, ipAB, α, β, backend, allocator)
+        tensorfree!(A_, allocator)
+        return C
+    end
+end
+function _blas_contract_B!(C, A, B, pB, pAB, ipAB, α, β, backend, allocator)
+    if isblascontractable(B, pB) && eltype(B) == eltype(C)
+        return _blas_contract_C!(C, A, _blas_matrix(B, pB), pAB, ipAB, α, β, backend, allocator)
+    else
+        B_ = tensoralloc_add(eltype(C), B, pB, false, Val(true), allocator)
+        Bnew = tensoradd!(SV(B_), B, pB, false, One(), Zero(), backend, allocator)
+        Bmat = _blas_matrix(Bnew, trivialpermutation(pB))
+        _blas_contract_C!(C, A, Bmat, pAB, ipAB, α, β, backend, allocator)
+        tensorfree!(B_, allocator)
+        return C
+    end
+end
+function _blas_contract_C!(C, A, B, pAB, ipAB, α, β, backend, allocator)
+    if isblasdestination(C, ipAB)
+        mul!(_blas_matrix(C, ipAB), A, B, α, β)
     else
+        TC = eltype(C)
         C_ = SV(tensoralloc_add(TC, C, ipAB, false, Val(true), allocator))
-        _unsafe_blas_contract!(
-            C_, A_, pA, B_, pB, trivialpermutation(ipAB′),
-            one(TC), zero(TC)
-        )
+        mul!(_blas_matrix(C_, trivialpermutation(ipAB)), A, B, one(TC), zero(TC))
         # `C` already exists, so only the permutation of `pAB` matters here and it can be
         # handed over in the canonical `{N,0}` partition
         tensoradd!(C, C_, (pAB, ()), false, α, β, backend, allocator)
         tensorfree!(C_.parent, allocator)
     end
-    flagA || tensorfree!(A_.parent, allocator)
-    flagB || tensorfree!(B_.parent, allocator)
     return C
 end
 
-# perform the actual contraction, assuming it can be done as matrix multiplication by simply
-# reshaping without any further allocations
-function _unsafe_blas_contract!(
-        C::StridedView{T},
-        A::StridedView{T}, pA,
-        B::StridedView{T}, pB,
-        pAB::IndexTuple, α, β
-    ) where {T <: BlasFloat}
-    sizeA = size(A)
-    sizeB = size(B)
-    csizeA = TupleTools.getindices(sizeA, pA[2])
-    csizeB = TupleTools.getindices(sizeB, pB[1])
-    osizeA = TupleTools.getindices(sizeA, pA[1])
-    osizeB = TupleTools.getindices(sizeB, pB[2])
-
-    mul!(
-        sreshape(permutedims(C, pAB), (prod(osizeA), prod(osizeB))),
-        sreshape(permutedims(A, linearize(pA)), (prod(osizeA), prod(csizeA))),
-        sreshape(permutedims(B, linearize(pB)), (prod(csizeB), prod(osizeB))),
-        α, β
-    )
-
-    return C
-end
-
-function makeblascontractable(A, pA, TC, backend, allocator)
-    flagA = isblascontractable(A, pA) && eltype(A) == TC
-    if !flagA
-        A_ = tensoralloc_add(TC, A, pA, false, Val(true), allocator)
-        Anew = SV(A_)
-        Anew = tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator)
-        pAnew = trivialpermutation(pA)
-    else
-        Anew = A
-        pAnew = pA
-    end
-    return Anew, pAnew, flagA
+# reshape `A` into the matrix with the indices `p[1]` as rows and `p[2]` as columns, assuming
+# this is possible without any further allocations
+function _blas_matrix(A::StridedView, p::Index2Tuple)
+    sz = size(A)
+    m = prod(TupleTools.getindices(sz, p[1]))
+    n = prod(TupleTools.getindices(sz, p[2]))
+    return sreshape(permutedims(A, linearize(p)), (m, n))
 end
 
 function isblascontractable(A::StridedView, p::Index2Tuple)

src/implementation/indices.jl

@@ -81,16 +81,13 @@ function trace_indices(IA::NTuple{NA, Any}, IC::NTuple{NC, Any}) where {NA, NC}
     return (indCinA, ()), (cindA1, cindA2)
 end
 
-function contract_indices(IA, IB, IC)
-    return contract_indices(tuple(IA...), tuple(IB...), tuple(IC...))
-end
+contract_indices(IA, IB, IC) = contract_indices(tuple(IA...), tuple(IB...), tuple(IC...))
 function contract_indices(
         IA::NTuple{NA, Any}, IB::NTuple{NB, Any}, IC::NTuple{NC, Any}
     ) where {NA, NB, NC}
     # labels
     IAB = (IA..., IB...)
-    isodd(length(IAB) - length(IC)) &&
-        throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+    isodd(length(IAB) - length(IC)) && _throw_invalid_contraction(IA, IB, IC)
     Icontract = tunique(tsetdiff(IAB, IC))
     IopenA = tsetdiff(IA, Icontract)
     IopenB = tsetdiff(IB, Icontract)
@@ -107,12 +104,58 @@ function contract_indices(
     pAB = (indCinoAB, ())
 
     if !isperm(linearize(pA)) || !isperm(linearize(pB)) || !isperm(indCinoAB)
-        throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+        _throw_invalid_contraction(IA, IB, IC)
     end
 
     return pA, pB, pAB
 end
 
+# Implementation for labels stored in vectors, where the number of labels is only known at
+# runtime. The permutations are computed as vectors in a type stable manner, and only then
+# converted to tuples. This avoids the type instabilities of the tuple implementation, as well
+# as compiling it for every combination of lengths.
+# note: inlined such that callers with a function barrier, such as `ncon`, can pass the
+# resulting tuples on without boxing them
+@inline function contract_indices(IA::AbstractVector, IB::AbstractVector, IC::AbstractVector)
+    # split the indices of `A` into open and contracted ones
+    oindA, cindA, cindB = Int[], Int[], Int[]
+    for (i, l) in enumerate(IA)
+        j = _findfirst(isequal(l), IB)
+        if j == 0
+            push!(oindA, i)
+        else
+            push!(cindA, i)
+            push!(cindB, j)
+        end
+    end
+    oindB = findall(l -> !(l in IA), IB)
+    # position of the indices of `C` in the open indices of `A` and `B`
+    indCinoAB = map(IC) do l
+        i = findfirst(i -> isequal(IA[i], l), oindA)
+        isnothing(i) || return i
+        j = findfirst(j -> isequal(IB[j], l), oindB)
+        return isnothing(j) ? 0 : length(oindA) + j
+    end
+    # every index of `B` and `C` should be accounted for exactly once
+    (
+        length(cindB) + length(oindB) == length(IB) && allunique(cindB) &&
+            length(oindA) + length(oindB) == length(IC) && allunique(indCinoAB) &&
+            !(0 in indCinoAB)
+    ) || throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+
+    pA = (_totuple(oindA), _totuple(cindA))
+    pB = (_totuple(cindB), _totuple(oindB))
+    pAB = (_totuple(indCinoAB), ())
+    return pA, pB, pAB
+end
+
+# note: `ntuple` is much faster than `Tuple` or splatting for short lists, and the separate
+# function avoids inferring the resulting union of tuple types in the caller
+@noinline _totuple(p::Vector{Int}) = ntuple(k -> p[k], length(p))
+
+@noinline _throw_invalid_contraction(IA, IB, IC) =
+    throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+
 #-------------------------------------------------------------------------------------------
 # Convert indices to einsum labels: useful for package extensions / add-ons (e.g. TBLIS)
 #-------------------------------------------------------------------------------------------

src/implementation/ncon.jl

@@ -39,43 +39,92 @@ function ncon(
         end
     end
 
-    tensors, network = resolve_traces(tensors, network)
-    tree = order === nothing ? ncontree(network) : indexordertree(network, order)
+    tensors′, network′ = resolve_traces(tensors, network)
+    # extract the keyword arguments once, rather than for every pairwise contraction
+    backend, allocator = _ncon_kwargs(; kwargs...)
+    conjlist′ = _ncon_conjlist(conjlist)
+    return _ncon(tensors′, network′, conjlist′, output′, order, backend, allocator)
+end
 
-    allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator()
+function _ncon(tensors, network, conjlist, output, order, backend, allocator)
+    steps = order === nothing ? _ncontree_steps(network) :
+        _indexordertree_steps(network, order)
     cp = allocator_checkpoint!(allocator)
-
-    A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
-    B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
-    IC = tuple(output′...)
-    C = tensorcontract(IC, A, IA, conjA, B, IB, conjB; kwargs...)
-    tree[1] isa Int || tensorfree!(A, allocator)
-    tree[2] isa Int || tensorfree!(B, allocator)
+    n = length(tensors)
+    i, j = steps[end]
+    A, IA, conjA = contracttree(tensors, network, conjlist, steps, i, backend, allocator)
+    B, IB, conjB = contracttree(tensors, network, conjlist, steps, j, backend, allocator)
+    C = _ncon_contract(A, IA, conjA, B, IB, conjB, output, Val(false), backend, allocator)
+    i > n && tensorfree!(A, allocator)
+    j > n && tensorfree!(B, allocator)
     allocator_reset!(allocator, cp)
-    return length(IC) == 0 ? tensorscalar(C) : C
+    return isempty(output) ? tensorscalar(C) : C
+end
+
+# Compute the (partial) result `k` of the contraction tree `steps` (see `_mergetrees!`).
+# The types and number of indices of the intermediate tensors are only known at runtime, but
+# the labels are kept as `Vector{Int}` such that all index computations are type stable, and
+# every pairwise contraction goes through a single dynamic dispatch, in `_ncon_contract`.
+function contracttree(tensors, network, conjlist, steps, k::Int, backend, allocator)
+    n = length(tensors)
+    k <= n && return tensors[k], network[k], conjlist[k]
+    i, j = steps[k - n]
+    A, IA, conjA = contracttree(tensors, network, conjlist, steps, i, backend, allocator)
+    B, IB, conjB = contracttree(tensors, network, conjlist, steps, j, backend, allocator)
+    IC = _ncon_symdiff(IA, IB)
+    C = _ncon_contract(A, IA, conjA, B, IB, conjB, IC, Val(true), backend, allocator)
+    i > n && tensorfree!(A, allocator)
+    j > n && tensorfree!(B, allocator)
+    return C, IC, false
 end
 
-function contracttree(tensors, network, conjlist, tree; kwargs...)
-    @nospecialize
-    if tree isa Int
-        return tensors[tree], tuple(network[tree]...), (conjlist[tree])
+function _ncon_kwargs(; kwargs...)
+    args = _kwargs2args(; kwargs...)
+    return length(args) == 2 ? args :
+        length(args) == 1 ? (args[1], DefaultAllocator()) :
+        (DefaultBackend(), DefaultAllocator())
+end
+
+# labels of the result of contracting tensors with labels `IA` and `IB`, i.e. all labels that
+# appear in only one of them, in order of appearance
+# note: this is much faster than `Base.symdiff`, which constructs a `Set`
+function _ncon_symdiff(IA::Vector{Int}, IB::Vector{Int})
+    IC = Int[]
+    for i in IA
+        i in IB || push!(IC, i)
     end
-    A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
-    B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
-    IC = tuple(symdiff(IA, IB)...)
-    pA, pB, _ = contract_indices(IA, IB, IC)
-    # `IC` lists the open indices of `A` before those of `B`, as in the product
-    pAB = trivialpermutation(length(pA[1]), length(pB[2]))
-    TC = promote_contract(scalartype(A), scalartype(B))
-    allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator()
-    backend = haskey(kwargs, :backend) ? kwargs[:backend] : DefaultBackend()
-    C = tensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, Val(true), allocator)
-    C = tensorcontract!(
-        C, A, pA, conjA, B, pB, conjB, pAB, One(), Zero(), backend, allocator
+    for i in IB
+        i in IA || push!(IC, i)
+    end
+    return IC
+end
+
+function _ncon_contract(
+        A, IA::Vector{Int}, conjA::Bool, B, IB::Vector{Int}, conjB::Bool, IC::Vector{Int},
+        istemp::Val, backend, allocator
+    )
+    # only `_ncon_contract!` should be specialized on the types of the tensors
+    @nospecialize A B
+    pA, pB, pAB = contract_indices(IA, IB, IC)
+    # function barrier: the number of indices of `A`, `B` and `C` is only known at runtime
+    # note: the tuples are passed separately rather than as `Index2Tuple`s to avoid boxing
+    return _ncon_contract!(
+        A, pA[1], pA[2], conjA, B, pB[1], pB[2], conjB, pAB[1], istemp, backend, allocator
     )
-    tree[1] isa Int || tensorfree!(A, allocator)
-    tree[2] isa Int || tensorfree!(B, allocator)
-    return C, IC, false
+end
+function _ncon_contract!(
+        A, oindA, cindA, conjA, B, cindB, oindB, conjB, indCinoAB, istemp::Val,
+        backend, allocator
+    )
+    pA = (oindA, cindA)
+    pB = (cindB, oindB)
+    # intermediates list the open indices of `A` before those of `B`, so partition them as
+    # `(open A; open B)`, which avoids repartitioning copies for e.g. symmetric tensors
+    pAB = istemp === Val(true) ? trivialpermutation(length(oindA), length(oindB)) :
+        (indCinoAB, ())
+    TC = promote_contract(scalartype(A), scalartype(B))
+    C = tensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, istemp, allocator)
+    return tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, One(), Zero(), backend, allocator)
 end
 
 function nconoutput(network, output)
@@ -91,5 +140,7 @@ function nconoutput(network, output)
 
     issetequal(output, outputindices) ||
         throw(ArgumentError("invalid NCON network: $network -> $output"))
-    return output
+    return collect(Int, output)
 end
+
+_ncon_conjlist(conjlist) = collect(Bool, conjlist)