diff --git a/docs/src/changelog.md b/docs/src/changelog.md index c8f910b3..399232cc 100644 --- a/docs/src/changelog.md +++ b/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 diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl index 1ca4c721..311f0345 100644 --- a/ext/TensorOperationsChainRulesCoreExt.jl +++ b/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._ncon_network(args...) +@non_differentiable TensorOperations._ncon_labelvector(args...) +@non_differentiable TensorOperations._ncon_conjlist(args...) +@non_differentiable TensorOperations._ncon_symdiff(args...) end diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl index 499dbe30..375d6bb3 100644 --- a/src/TensorOperations.jl +++ b/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") diff --git a/src/implementation/blascontract.jl b/src/implementation/blascontract.jl index 72bb9278..c0e1325e 100644 --- a/src/implementation/blascontract.jl +++ b/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) diff --git a/src/implementation/indices.jl b/src/implementation/indices.jl index 06ddb6c0..bafff825 100644 --- a/src/implementation/indices.jl +++ b/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) #------------------------------------------------------------------------------------------- diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl index 79f06df3..03d688a2 100644 --- a/src/implementation/ncon.jl +++ b/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) diff --git a/src/indexnotation/indexordertree.jl b/src/indexnotation/indexordertree.jl index c1ad07bc..2a87a6af 100644 --- a/src/indexnotation/indexordertree.jl +++ b/src/indexnotation/indexordertree.jl @@ -1,63 +1,33 @@ # verify that every conraction index appears exactly twice, and any other index exactly once function _verify_indexorder(network, indexorder) - allindices = Vector{Any}() - for ind in network - append!(allindices, ind) - end - while length(allindices) > 0 - i = pop!(allindices) - if i ∈ indexorder # contraction index - k = findfirst(isequal(i), allindices) - k === nothing && return false - l = findnext(isequal(i), allindices, k + 1) - l !== nothing && return false - deleteat!(allindices, k) - else - findfirst(isequal(i), allindices) === nothing || return false + for indices in network + for i in indices + n = sum(ind -> count(isequal(i), ind), network) + n == (i in indexorder ? 2 : 1) || return false end end return true end function indexordertree(network, indexorder) - _verify_indexorder(network, indexorder) || - throw(ArgumentError("not a valid contraction tree with contraction indices $indexorder")) - contractionindices = Vector{Vector{Any}}(undef, length(network)) - for k in 1:length(network) - indices = network[k] - # trace indices have already been removed, remove open indices by filtering on positive values - contractionindices[k] = filter(in(indexorder), indices) - end - partialtrees = collect(Any, 1:length(network)) - return _indexordertree!(partialtrees, contractionindices, indexorder) + return _nestedtree(_indexordertree_steps(network, indexorder), length(network)) end -function _indexordertree!(partialtrees, contractionindices, indexorder) - while length(indexorder) > 0 - c = indexorder[1] - i1 = findfirst(x -> in(c, x), contractionindices) - if i1 === nothing - indexorder = indexorder[2:end] - else - i2 = findnext(x -> in(c, x), contractionindices, i1 + 1) - @assert i2 !== nothing - cs = intersect(contractionindices[i1], contractionindices[i2]) - indexorder = setdiff(indexorder, cs) - newindices = unique2(vcat(contractionindices[i1], contractionindices[i2])) - newtree = Any[partialtrees[i1], partialtrees[i2]] - partialtrees[i1] = newtree - deleteat!(partialtrees, i2) - contractionindices[i1] = newindices - deleteat!(contractionindices, i2) - end - end - if length(partialtrees) > 1 - @assert all(isempty, contractionindices) # disconnected network - while length(partialtrees) > 1 - partialtrees[end - 1] = Any[partialtrees[end - 1], partialtrees[end]] - pop!(partialtrees) - pop!(contractionindices) - end +# Determine the contraction order where the indices are contracted in the order specified by +# `indexorder`, as a list of pairwise contractions (see `_mergetrees!`). +function _indexordertree_steps(network, indexorder) + _verify_indexorder(network, indexorder) || + throw(ArgumentError("not a valid contraction tree with contraction indices $indexorder")) + steps = NTuple{2, Int}[] + partialtrees = collect(1:length(network)) + group = collect(1:length(network)) + for c in indexorder + k1 = findfirst(indices -> c in indices, network) + isnothing(k1) && continue + # trace indices have already been removed, so `c` appears in another tensor + k2 = findnext(indices -> c in indices, network, k1 + 1) + isnothing(k2) && continue + _mergetrees!(steps, partialtrees, group, k1, k2) end - return partialtrees[1] + return _finalizetree!(steps, partialtrees, group) end diff --git a/src/indexnotation/ncontree.jl b/src/indexnotation/ncontree.jl index ed8f6f88..be93f042 100644 --- a/src/indexnotation/ncontree.jl +++ b/src/indexnotation/ncontree.jl @@ -5,63 +5,107 @@ check_nconstyle(network) = _check_nconstyle(network, Val(false)) function _check_nconstyle(network, ::Val{check}) where {check} allindices = Vector{Int}() for ind in network - all(i -> isa(i, Integer), ind) || return check ? false : - throw(IndexError("All indices must be integers")) - append!(allindices, ind) + for i in ind + isa(i, Integer) || return check ? false : + throw(IndexError("All indices must be integers")) + push!(allindices, i) + end end - while length(allindices) > 0 - i = pop!(allindices) + sort!(allindices) + j = 1 + while j <= length(allindices) + i = allindices[j] + # number of times that `i` appears + n = 1 + while j + n <= length(allindices) && allindices[j + n] == i + n += 1 + end if i > 0 # positive labels represent contractions or traces and should appear twice - k = findfirst(isequal(i), allindices) - isnothing(k) && return check ? false : + n == 1 && return check ? false : throw(IndexError(lazy"Index $i appears only once in the network")) - l = findnext(isequal(i), allindices, k + 1) - !isnothing(l) && return check ? false : + n > 2 && return check ? false : throw(IndexError(lazy"Index $i appears more than twice in the network")) - deleteat!(allindices, k) elseif i < 0 # negative labels represent open indices and should appear once - isnothing(findfirst(isequal(i), allindices)) || return check ? false : + n > 1 && return check ? false : throw(IndexError(lazy"Index $i appears more than once in the network")) else # i == 0 return check ? false : throw(IndexError("Index 0 is not allowed in the network")) end + j += n end return check ? true : nothing end -function ncontree(network) - contractionindices = Vector{Vector{Int}}(undef, length(network)) - for k in 1:length(network) - indices = network[k] - # trace indices have already been removed, remove open indices by filtering on positive values - contractionindices[k] = filter(i -> i > 0, indices) +ncontree(network) = _nestedtree(_ncontree_steps(network), length(network)) + +# Determine the contraction order where the indices are contracted in increasing order, as a +# list of pairwise contractions (see `_mergetrees!`). +function _ncontree_steps(network) + # collect the positive labels, i.e. the contraction indices, and the tensors they belong to + contractionindices = Int[] + owners = Int[] + for (k, indices) in enumerate(network) + for i in indices + i > 0 || continue + push!(contractionindices, i) + push!(owners, k) + end + end + + # contract the indices in increasing order + steps = NTuple{2, Int}[] + partialtrees = collect(1:length(network)) + group = collect(1:length(network)) + for j in sortperm(contractionindices) + k = findnext(==(contractionindices[j]), contractionindices, j + 1) + # each index appears twice, only consider the first occurrence + isnothing(k) || _mergetrees!(steps, partialtrees, group, owners[j], owners[k]) end - partialtrees = collect(Any, 1:length(network)) - return _ncontree!(partialtrees, contractionindices) + return _finalizetree!(steps, partialtrees, group) +end + +# A contraction tree for a network of `n` tensors is represented as a list of pairwise +# contractions `steps`, where `steps[s] = (i, j)` contracts the (partial) results `i` and `j`. +# Here, results `1:n` are the tensors themselves, and result `n + s` is the result of `steps[s]`. +# The partial tree containing tensor `k` is `partialtrees[group[k]]`. + +# Merge the partial trees containing tensors `k1` and `k2`. The merged tree replaces the one +# containing the tensor that comes first, and nothing happens if both tensors are already in +# the same partial tree. +function _mergetrees!( + steps::Vector{NTuple{2, Int}}, partialtrees::Vector{Int}, group::Vector{Int}, + k1::Int, k2::Int + ) + g1, g2 = minmax(group[k1], group[k2]) + g1 == g2 && return steps + push!(steps, (partialtrees[g1], partialtrees[g2])) + partialtrees[g1] = length(group) + length(steps) + replace!(group, g2 => g1) + return steps end -function _ncontree!(partialtrees, contractionindices) - if length(partialtrees) == 1 - return partialtrees[1] +# Combine the remaining partial trees, which are disconnected and thus require outer products, +# starting from the end. +function _finalizetree!( + steps::Vector{NTuple{2, Int}}, partialtrees::Vector{Int}, group::Vector{Int} + ) + n = length(group) + remaining = unique!(group) + tree = partialtrees[remaining[end]] + for g in Iterators.drop(Iterators.reverse(remaining), 1) + push!(steps, (partialtrees[g], tree)) + tree = n + length(steps) end - if all(isempty, contractionindices) # disconnected network - partialtrees[end - 1] = Any[partialtrees[end - 1], partialtrees[end]] - pop!(partialtrees) - pop!(contractionindices) - else - let firstind = minimum(vcat(contractionindices...)) - i1 = findfirst(x -> in(firstind, x), contractionindices) - i2 = findnext(x -> in(firstind, x), contractionindices, i1 + 1) - @assert i1 !== nothing && i2 !== nothing - newindices = unique2(vcat(contractionindices[i1], contractionindices[i2])) - newtree = Any[partialtrees[i1], partialtrees[i2]] - partialtrees[i1] = newtree - deleteat!(partialtrees, i2) - contractionindices[i1] = newindices - deleteat!(contractionindices, i2) - end + return steps +end + +# convert a list of pairwise contractions to a nested contraction tree +function _nestedtree(steps::Vector{NTuple{2, Int}}, n::Int) + trees = collect(Any, 1:n) + for (i, j) in steps + push!(trees, Any[trees[i], trees[j]]) end - return _ncontree!(partialtrees, contractionindices) + return trees[end] end """ @@ -97,10 +141,25 @@ function nconindexcompletion(ex) end end +# Resolve the traces in the network, and canonicalize the types of the tensor list and the +# network, such that the remainder of `ncon` is type stable and only compiled once. function resolve_traces(tensors, network) - transformed = map(zip(tensors, network)) do (A, IA) - IC = unique2(IA) - return length(IC) == length(IA) ? (A, IA) : (tensortrace(IC, A, IA), IC) + network′ = _ncon_network(network) + all(k -> length(network[k]) == length(network′[k]), eachindex(network)) && + return _anyvector(tensors), network′ + tensors′ = map(tensors, network, network′) do A, IA, IC + return length(IA) == length(IC) ? A : tensortrace(IC, A, _ncon_labelvector(IA)) end - return first.(transformed), last.(transformed) + return _anyvector(tensors′), network′ end + +# labels after resolving the traces +_ncon_network(network) = + Vector{Int}[_ncon_labelvector(allunique(IA) ? IA : unique2(IA)) for IA in network] + +# canonicalize any list of labels, e.g. a `Tuple` or a row matrix like `[1 2 -1]` +_ncon_labelvector(IA) = vec(collect(Int, IA)) + +# note: avoiding `collect(Any, xs)` and mutation for compatibility with Zygote +_anyvector(xs::AbstractVector) = Vector{Any}(xs) +_anyvector(xs::Tuple) = Any[xs...] diff --git a/test/auxiliary.jl b/test/auxiliary.jl index fcfa95af..e102ca03 100644 --- a/test/auxiliary.jl +++ b/test/auxiliary.jl @@ -336,5 +336,39 @@ end ) @test contract_indices(contract_labels(pA, pB, pAB)...) == (pA, pB, pAB) end + + # labels stored in vectors + @test contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :d]) == + (((1, 2), (3,)), ((1,), (2,)), ((1, 2, 3), ())) + @test contract_indices([:c, :a, :b], [:c, :d], [:a, :b, :d]) == + (((2, 3), (1,)), ((1,), (2,)), ((1, 2, 3), ())) + @test contract_indices([1, 2], [3, 4], [3, 1, 4, 2]) == + (((1, 2), ()), ((), (1, 2)), ((3, 1, 4, 2), ())) + @test contract_indices([1, 2], [2, 1], Int[]) == + (((), (1, 2)), ((2, 1), ()), ((), ())) + @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :d, :e]) + @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :e]) + @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d, :e], [:a, :b, :d]) + @test_throws IndexError contract_indices([:a, :b, :c], [:c, :c], [:a, :b]) + + # agrees with the tuple implementation, including for invalid patterns + result(f) = try + f() + catch e + e isa IndexError || rethrow() + IndexError + end + for _ in 1:1000 + c = collect(1:rand(0:3)) + IA = shuffle([c; 100 .+ (1:rand(0:3))]) + IB = shuffle([c; 200 .+ (1:rand(0:3))]) + IC = shuffle(filter(>(99), [IA; IB])) + r = rand(1:4) + r == 1 && !isempty(IC) && (IC[1] = 999) + r == 2 && !isempty(IA) && (IA[1] = 777) + r == 3 && length(IB) > 1 && (IB[2] = IB[1]) + @test result(() -> contract_indices(IA, IB, IC)) == + result(() -> contract_indices(Tuple(IA), Tuple(IB), Tuple(IC))) + end end end diff --git a/test/tensor.jl b/test/tensor.jl index 2ac3f71e..11ed2473 100644 --- a/test/tensor.jl +++ b/test/tensor.jl @@ -543,40 +543,24 @@ end end @testset "methods for StridedBLAS" begin - using TensorOperations: isblascontractable, isblasdestination, makeblascontractable, - StridedBLAS, DefaultAllocator + using TensorOperations: isblascontractable using Strided - backend = StridedBLAS() - allocator = DefaultAllocator() A = StridedView(rand(ComplexF64, 4, 4, 4, 4)) p = ((1, 2), (3, 4)) @test isblascontractable(A, p) - Anew, pnew, flag = makeblascontractable(A, p, ComplexF64, backend, allocator) - @test isblascontractable(Anew, pnew) - @test Anew === A - @test !isblascontractable(conj(A), p) - Anew, pnew, flag = makeblascontractable(A, p, ComplexF64, backend, allocator) - @test isblascontractable(Anew, pnew) for p in (((2, 1), (3, 4)), ((1,), (3, 2, 4)), ((2, 1, 4), (3,))), op in (identity, conj) @test !isblascontractable(op(A), p) - Anew, pnew, flag = makeblascontractable( - op(A), p, ComplexF64, backend, allocator - ) - @test isblascontractable(Anew, pnew) end vA = view(A, 1:2, 1:2, 1:2, 1:2) p = ((1, 2), (3, 4)) @test !isblascontractable(vA, p) - Anew, pnew, flag = makeblascontractable(vA, p, ComplexF64, backend, allocator) - @test !flag - @test isblascontractable(Anew, pnew) pA = permutedims(A, (3, 4, 1, 2)) p = ((1, 2), (3, 4))