Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/spaces/homspace.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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 &&
Expand Down Expand Up @@ -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))...)

Expand Down
11 changes: 9 additions & 2 deletions src/tensors/indexmanipulations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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₂)
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down
61 changes: 35 additions & 26 deletions src/tensors/tensoroperations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand All @@ -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!(
Expand Down
Loading