update dependencies, update to internal API of TensorOperations 3; sw…

…itch to `@spawn` for multithreaded abelian tensors
Jutho · Jun 3, 2020 · 4496a2a · 4496a2a
1 parent 0af6885
commit 4496a2a
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 90 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "TensorKit"
 uuid = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 authors = ["Jutho Haegeman"]
-version = "0.4.3"
+version = "0.5.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -13,13 +13,13 @@ TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 
 [compat]
-Strided = "0.3.3"
+Strided = "1.0.1"
 TupleTools = "1.1"
 HalfIntegers = "1"
 WignerSymbols = "1"
-TensorOperations = "2.1"
+TensorOperations = "3"
 LRUCache = "1.0.2"
-julia = "1.4, ~1.0.6"
+julia = "1.4"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl
@@ -1,64 +1,3 @@
-# import TensorOperations: checked_similar_from_indices, scalar, isblascontractable, add!, contract!, unsafe_contract!
-
-# Add support for TensorOperations.jl
-function TensorOperations.memsize(t::TensorMap)
- s = 0
- for (c,b) in blocks(t)
- s += sizeof(b)
- end
- return s
-end
-TensorOperations.memsize(t::AdjointTensorMap) = TensorOperations.memsize(t')
-
-function checked_similar_from_indices(tC, ::Type{T}, p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
- t::AbstractTensorMap{S}) where {T,S<:IndexSpace,N₁,N₂}
-
- cod = ProductSpace{S,N₁}(space.(Ref(t), p1))
- dom = ProductSpace{S,N₂}(dual.(space.(Ref(t), p2)))
- A = similarstoragetype(t, T)
-
- if tC !== nothing && (tC isa TensorMap{S,N₁,N₂}) && eltype(tC) == T &&
- codomain(tC) == cod && domain(tC) == dom && storagetype(tC) == A
- G = sectortype(S)
- if G === Trivial
- return tC::TensorMap{S,N₁,N₂,G,A,Nothing,Nothing}
- else
- F₁ = fusiontreetype(G, StaticLength(N₁))
- F₂ = fusiontreetype(G, StaticLength(N₂))
- return tC::TensorMap{S,N₁,N₂,G,SectorDict{G,A},F₁,F₂}
- end
- else
- tC = similar(t, T, cod←dom)
- return tC
- end
-end
-
-function checked_similar_from_indices(tC, ::Type{T}, oindA::IndexTuple, oindB::IndexTuple,
- p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, tA::AbstractTensorMap{S},
- tB::AbstractTensorMap{S}) where {T, S<:IndexSpace,N₁,N₂}
-
- spaces = (space.(Ref(tA), oindA)..., space.(Ref(tB), oindB)...)
- cod = ProductSpace{S,N₁}(getindex.(Ref(spaces), p1))
- dom = ProductSpace{S,N₂}(dual.(getindex.(Ref(spaces), p2)))
- A = similarstoragetype(tA, T)
- @assert A === similarstoragetype(tB, T)
-
- if tC !== nothing && (tC isa TensorMap{S,N₁,N₂}) && eltype(tC) == T &&
- codomain(tC) == cod && domain(tC) == dom && storagetype(tC) == A
- G = sectortype(S)
- if G === Trivial
- return tC::TensorMap{S,N₁,N₂,G,A,Nothing,Nothing}
- else
- F₁ = fusiontreetype(G, StaticLength(N₁))
- F₂ = fusiontreetype(G, StaticLength(N₂))
- return tC::TensorMap{S,N₁,N₂,G,SectorDict{G,A},F₁,F₂}
- end
- else
- tC = similar(tA, T, cod←dom)
- return tC
- end
-end
-
 function cached_permute(sym::Symbol, t::TensorMap{S},
  p1::IndexTuple{N₁}, p2::IndexTuple{N₂}=()) where {S,N₁,N₂}
  cod = ProductSpace{S,N₁}(map(n->space(t, n), p1))
@@ -127,18 +66,9 @@ function add!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S,N₁
  pdata = (p1..., p2...)
  axpby!(α, permutedims(tsrc[], pdata), β, tdst[])
  elseif FusionStyle(G) isa Abelian && BraidingStyle(G) isa SymmetricBraiding
- K = Threads.nthreads()
- if K > 1
- let iterator = fusiontrees(tsrc)
- Threads.@threads for k = 1:K
- counter = 0
- for (f1,f2) in iterator
- counter += 1
- if mod1(counter, K) == k
- _addabelianblock!(α, tsrc, β, tdst, p1, p2, f1, f2)
- end
- end
- end
+ if Threads.nthreads() > 1
+ Threads.@sync for (f1,f2) in fusiontrees(tsrc)
+ Threads.@spawn _addabelianblock!(α, tsrc, β, tdst, p1, p2, f1, f2)
  end
  else # debugging is easier this way
  for (f1,f2) in fusiontrees(tsrc)
@@ -319,38 +249,65 @@ function _contract!(α, A::AbstractTensorMap{S}, B::AbstractTensorMap{S},
  return C
 end
 
-# # Compatibility layer for working with the `@tensor` macro from TensorOperations
-function TO.checked_similar_from_indices(C, T::Type,
- p1::IndexTuple, p2::IndexTuple, A::AbstractTensorMap, CA::Symbol = :N)
+# Add support for cache and API (`@tensor` macro & friends) from TensorOperations.jl:
+# compatibility layer
+function TensorOperations.memsize(t::TensorMap)
+ s = 0
+ for (c,b) in blocks(t)
+ s += sizeof(b)
+ end
+ return s
+end
+TensorOperations.memsize(t::AdjointTensorMap) = TensorOperations.memsize(t')
+
+function TO.similarstructure_from_indices(T::Type, p1::IndexTuple, p2::IndexTuple,
+ A::AbstractTensorMap, CA::Symbol = :N)
  if CA == :N
- checked_similar_from_indices(C, T, p1, p2, A)
+ _similarstructure_from_indices(T, p1, p2, A)
  else
  p1 = adjointtensorindices(A, p1)
  p2 = adjointtensorindices(A, p2)
- checked_similar_from_indices(C, T, p1, p2, adjoint(A))
+ _similarstructure_from_indices(T, p1, p2, adjoint(A))
  end
 end
 
-function TO.checked_similar_from_indices(C, T::Type,
- poA::IndexTuple, poB::IndexTuple,
- p1::IndexTuple, p2::IndexTuple,
- A::AbstractTensorMap, B::AbstractTensorMap, CA::Symbol = :N, CB::Symbol = :N)
+function TO.similarstructure_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple,
+  p1::IndexTuple, p2::IndexTuple,
+  A::AbstractTensorMap, B::AbstractTensorMap,
+   CA::Symbol = :N, CB::Symbol = :N)
 
  if CA == :N && CB == :N
- checked_similar_from_indices(C, T, poA, poB, p1, p2, A, B)
+ _similarstructure_from_indices(T, poA, poB, p1, p2, A, B)
  elseif CA == :C && CB == :N
  poA = adjointtensorindices(A, poA)
- checked_similar_from_indices(C, T, poA, poB, p1, p2, adjoint(A), B)
+ _similarstructure_from_indices(T, poA, poB, p1, p2, adjoint(A), B)
  elseif CA == :N && CB == :C
  poB = adjointtensorindices(B, poB)
- checked_similar_from_indices(C, T, poA, poB, p1, p2, A, adjoint(B))
+ _similarstructure_from_indices(T, poA, poB, p1, p2, A, adjoint(B))
  else
  poA = adjointtensorindices(A, poA)
  poB = adjointtensorindices(B, poB)
- checked_similar_from_indices(C, T, poA, poB, p1, p2, adjoint(A), adjoint(B))
+ _similarstructure_from_indices(T, poA, poB, p1, p2, adjoint(A), adjoint(B))
  end
 end
 
+function _similarstructure_from_indices(::Type{T}, p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
+ t::AbstractTensorMap{S}) where {T,S<:IndexSpace,N₁,N₂}
+
+ cod = ProductSpace{S,N₁}(space.(Ref(t), p1))
+ dom = ProductSpace{S,N₂}(dual.(space.(Ref(t), p2)))
+ return dom→cod
+end
+function _similarstructure_from_indices(::Type{T}, oindA::IndexTuple, oindB::IndexTuple,
+ p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
+ tA::AbstractTensorMap{S}, tB::AbstractTensorMap{S}) where {T, S<:IndexSpace,N₁,N₂}
+
+ spaces = (space.(Ref(tA), oindA)..., space.(Ref(tB), oindB)...)
+ cod = ProductSpace{S,N₁}(getindex.(Ref(spaces), p1))
+ dom = ProductSpace{S,N₂}(dual.(getindex.(Ref(spaces), p2)))
+ return dom→cod
+end
+
 TO.scalar(t::AbstractTensorMap) = scalar(t)
 
 function TO.add!(α, tsrc::AbstractTensorMap{S}, CA::Symbol, β,