diff --git a/src/spaces/homspace.jl b/src/spaces/homspace.jl index 6e5153e83..8460f7578 100644 --- a/src/spaces/homspace.jl +++ b/src/spaces/homspace.jl @@ -91,7 +91,8 @@ sector structure of `W` (ignoring degeneracy dimensions). See also [`hasblock`](@ref), [`blockstructure`](@ref). """ -blocksectors(W::HomSpace) = sectorstructure(W).blocksectors +blocksectors(W::HomSpace) = + sectortype(W) === Trivial ? _blocksectors(W) : sectorstructure(W).blocksectors function _blocksectors(W::HomSpace) sectortype(W) === Trivial && @@ -130,7 +131,7 @@ hasblock(W::HomSpace, c::Sector) = c in blocksectors(W) Return the total dimension of a `HomSpace`, i.e. the number of linearly independent morphisms that can be constructed within this space. """ -dim(W::HomSpace) = degeneracystructure(W).totaldim +dim(W::HomSpace) = sectortype(W) === Trivial ? prod(dims(W)) : degeneracystructure(W).totaldim dims(W::HomSpace) = (dims(codomain(W))..., dims(domain(W))...) diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 17c23e026..7155bb4be 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -254,8 +254,7 @@ function permute( # general case tdst = similar(t, promote_permute(t), permute(space(t), p)) - levels = ntuple(identity, numind(t)) - return @inbounds braid!(tdst, t, p, levels, One(), Zero(), backend, allocator) + return @inbounds permute!(tdst, t, p, One(), Zero(), backend, allocator) end function permute(t::AdjointTensorMap, (p₁, p₂)::Index2Tuple; kwargs...) p₁′ = adjointtensorindices(t, p₂) @@ -306,6 +305,10 @@ See also [`braid`](@ref) for creating a new tensor. backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() ) @boundscheck spacecheck_transform(braid, tdst, tsrc, p, levels) + if has_array_view(tdst) && has_array_view(tsrc) + TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator) + return tdst + end levels1 = TupleTools.getindices(levels, codomainind(tsrc)) levels2 = TupleTools.getindices(levels, domainind(tsrc)) transformer = treebraider(tdst, tsrc, p, (levels1, levels2)) @@ -377,6 +380,10 @@ end backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator() ) @boundscheck spacecheck_transform(transpose, tdst, tsrc, p) + if has_array_view(tdst) && has_array_view(tsrc) + TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend, allocator) + return tdst + end transformer = treetransposer(tdst, tsrc, p) return @inbounds add_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) end diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 52f4eb417..147d27b16 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -47,6 +47,10 @@ function TO.tensoradd!( α::Number, β::Number, backend, allocator ) + if has_array_view(C) && has_array_view(A) + TO.tensoradd!(C[], A[], pA, conjA, α, β, backend, allocator) + return C + end if conjA A′ = adjoint(A) pA′ = adjointtensorindices(A, _canonicalize(pA, C)) @@ -320,31 +324,22 @@ function contract!( ) length(pA[2]) == length(pB[1]) || throw(IndexError("number of contracted indices does not match")) - N₁, N₂ = length(pA[1]), length(pB[2]) # find optimal contraction scheme by checking the following options: # - sorting the contracted inds of A or B to avoid permutations # - contracting B with A instead to avoid permutations + pA′, pB′, pA″, pB″, pAB′ = _contract_candidates(pA, pB, pAB) - qA = TupleTools.sortperm(pA[2]) - pA′ = Base.setindex(pA, TupleTools.getindices(pA[2], qA), 2) - pB′ = Base.setindex(pB, TupleTools.getindices(pB[1], qA), 1) - - qB = TupleTools.sortperm(pB[1]) - pA″ = Base.setindex(pA, TupleTools.getindices(pA[2], qB), 2) - pB″ = Base.setindex(pB, TupleTools.getindices(pB[1], qB), 1) + # dims are permutation-invariant, so compute them once here rather than in every memcost call + dA, dB, dC = dim(A), dim(B), dim(C) # keep order A en B, check possibilities for cind - memcost1 = TO.contract_memcost(C, A, pA′, B, pB′, pAB) - memcost2 = TO.contract_memcost(C, A, pA″, B, pB″, pAB) + memcost1 = _contract_memcost(dA, dB, dC, C, A, pA′, B, pB′, pAB) + memcost2 = _contract_memcost(dA, dB, dC, C, A, pA″, B, pB″, pAB) # reverse order A en B, check possibilities for cind - pAB′ = ( - map(n -> ifelse(n > N₁, n - N₁, n + N₂), pAB[1]), - map(n -> ifelse(n > N₁, n - N₁, n + N₂), pAB[2]), - ) - memcost3 = TO.contract_memcost(C, B, reverse(pB′), A, reverse(pA′), pAB′) - memcost4 = TO.contract_memcost(C, B, reverse(pB″), A, reverse(pA″), pAB′) + memcost3 = _contract_memcost(dB, dA, dC, C, B, reverse(pB′), A, reverse(pA′), pAB′) + memcost4 = _contract_memcost(dB, dA, dC, C, B, reverse(pB″), A, reverse(pA″), pAB′) return if min(memcost1, memcost2) <= min(memcost3, memcost4) if memcost1 <= memcost2 @@ -361,23 +356,37 @@ function contract!( end end -function TO.contract_memcost( - C::AbstractTensorMap, - A::AbstractTensorMap, pA::Index2Tuple, - B::AbstractTensorMap, pB::Index2Tuple, - pAB::Index2Tuple +# @noinline to avoid specialization +@noinline function _contract_candidates(pA::Index2Tuple, pB::Index2Tuple, pAB::Index2Tuple) + N₁, N₂ = length(pA[1]), length(pB[2]) + + qA = TupleTools.sortperm(pA[2]) + pA′ = Base.setindex(pA, TupleTools.getindices(pA[2], qA), 2) + pB′ = Base.setindex(pB, TupleTools.getindices(pB[1], qA), 1) + + qB = TupleTools.sortperm(pB[1]) + pA″ = Base.setindex(pA, TupleTools.getindices(pA[2], qB), 2) + pB″ = Base.setindex(pB, TupleTools.getindices(pB[1], qB), 1) + + pAB′ = ( + map(n -> ifelse(n > N₁, n - N₁, n + N₂), pAB[1]), + map(n -> ifelse(n > N₁, n - N₁, n + N₂), pAB[2]), ) + return pA′, pB′, pA″, pB″, pAB′ +end + +function _contract_memcost(dimA, dimB, dimC, C, A, pA, B, pB, pAB) ipAB = TO.oindABinC(pAB, pA, pB) - return dim(A) * (!TO.isblascontractable(A, pA) || eltype(A) !== eltype(C)) + - dim(B) * (!TO.isblascontractable(B, pB) || eltype(B) !== eltype(C)) + - dim(C) * !TO.isblasdestination(C, ipAB) + return dimA * (!TO.isblascontractable(A, pA) || scalartype(A) !== scalartype(C)) + + dimB * (!TO.isblascontractable(B, pB) || scalartype(B) !== scalartype(C)) + + dimC * !TO.isblasdestination(C, ipAB) end function TO.isblascontractable(A::AbstractTensorMap, pA::Index2Tuple) - return eltype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, pA) + return scalartype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, pA) end function TO.isblasdestination(A::AbstractTensorMap, ipAB::Index2Tuple) - return eltype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, ipAB) + return scalartype(A) <: LinearAlgebra.BlasFloat && has_shared_permute(A, ipAB) end function blas_contract!(