Improve contractions involving scalar-like tensors (#57)

* Add special code path for contractions involving scalar-like tensors

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
-version = "0.1.19"
+version = "0.1.20"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/NDTensors.jl

@@ -66,7 +66,7 @@ include("deprecated.jl")
 # A global timer used with TimerOutputs.jl
 #
 
-const GLOBAL_TIMER = TimerOutput()
+const timer = TimerOutput()
 
 #####################################
 # Optional TBLIS contraction backend

src/blocksparse/blocksparsetensor.jl

@@ -611,6 +611,7 @@ function permutedims!!(R::BlockSparseTensor{ElR,N},
                        T::BlockSparseTensor{ElT,N},
                        perm::NTuple{N,Int},
                        f::Function=(r,t)->t) where {ElR,ElT,N}
+  @timeit_debug timer "block sparse permutedims!!" begin
   bofsRR = blockoffsets(R)
   bofsT = blockoffsets(T)
 
@@ -643,6 +644,7 @@ function permutedims!!(R::BlockSparseTensor{ElR,N},
 
   permutedims!(R, T, perm, f)
   return R
+  end
 end
 
 # Version where it is known that R has the same blocks
@@ -781,6 +783,7 @@ function contraction_output(T1::TensorT1,
                             labelsT2,
                             labelsR) where {TensorT1<:BlockSparseTensor,
                                             TensorT2<:BlockSparseTensor}
+
   indsR = contract_inds(inds(T1),labelsT1,inds(T2),labelsT2,labelsR)
   TensorR = contraction_output_type(TensorT1,TensorT2,typeof(indsR))
   blockoffsetsR,contraction_plan = contract_blockoffsets(blockoffsets(T1),inds(T1),labelsT1,
@@ -795,9 +798,11 @@ function contract(T1::BlockSparseTensor{<:Any,N1},
                   T2::BlockSparseTensor{<:Any,N2},
                   labelsT2,
                   labelsR = contract_labels(labelsT1,labelsT2)) where {N1,N2}
+  @timeit_debug timer "Block sparse contract" begin
   R,contraction_plan = contraction_output(T1,labelsT1,T2,labelsT2,labelsR)
   R = contract!(R,labelsR,T1,labelsT1,T2,labelsT2,contraction_plan)
   return R
+  end
 end
 
 function contract!(R::BlockSparseTensor{ElR, NR},

src/blocksparse/combiner.jl

@@ -1,8 +1,7 @@
 
-function contract(T::BlockSparseTensor,
-                  labelsT,
-                  C::CombinerTensor,
-                  labelsC)
+function contract(T::BlockSparseTensor, labelsT,
+                  C::CombinerTensor, labelsC)
+  @timeit_debug timer "Block sparse (un)combiner" begin
   # Get the label marking the combined index
   # By convention the combined index is the first one
   # TODO: consider storing the location of the combined
@@ -40,16 +39,14 @@ function contract(T::BlockSparseTensor,
     Ruc = uncombine(T,indsRuc,cpos_in_labelsRc,blockperm(C),blockcomb(C))
     return Ruc
   end
+  end
 end
 
-contract(C::CombinerTensor,
-         labelsC,
-         T::BlockSparseTensor,
-         labelsT) = contract(T,labelsT,C,labelsC)
+contract(C::CombinerTensor, labelsC, T::BlockSparseTensor, labelsT) =
+  contract(T,labelsT,C,labelsC)
 
 # Special case when no indices are combined
-contract(T::BlockSparseTensor,
-         labelsT,
-         C::CombinerTensor{<:Any,0},
-         labelsC) = copy(T)
+# XXX: no copy
+contract(T::BlockSparseTensor, labelsT,
+         C::CombinerTensor{<:Any,0}, labelsC) = copy(T)
 

src/blocksparse/linearalgebra.jl

@@ -36,6 +36,7 @@ function LinearAlgebra.svd(T::BlockSparseMatrix{ElT};
 
   truncate = haskey(kwargs, :maxdim) || haskey(kwargs, :cutoff)
 
+  @timeit_debug timer "block sparse svd" begin
   Us = Vector{DenseTensor{ElT, 2}}(undef, nnzblocks(T))
   Ss = Vector{DiagTensor{real(ElT), 2}}(undef, nnzblocks(T))
   Vs = Vector{DenseTensor{ElT, 2}}(undef, nnzblocks(T))
@@ -180,6 +181,7 @@ function LinearAlgebra.svd(T::BlockSparseMatrix{ElT};
   end
 
   return U,S,V,Spectrum(d,truncerr)
+  end # @timeit_debug
 end
 
 _eigen_eltypes(T::Hermitian{ElT,<:BlockSparseMatrix{ElT}}) where {ElT} = real(ElT), ElT

src/combiner.jl

@@ -17,11 +17,11 @@ data(::Combiner) = error("Combiner storage has no data")
 blockperm(C::Combiner) = C.perm
 blockcomb(C::Combiner) = C.comb
 
-Base.eltype(::Type{<:Combiner}) = Number
+eltype(::Type{<:Combiner}) = Number
 
-Base.eltype(::Combiner) = eltype(Combiner)
+eltype(::Combiner) = eltype(Combiner)
 
-Base.promote_rule(::Type{<:Combiner},
+promote_rule(::Type{<:Combiner},
                   StorageT::Type{<:Dense}) = StorageT
 
 #
@@ -36,7 +36,7 @@ uncombinedinds(T::CombinerTensor) = popfirst(inds(T))
 blockperm(C::CombinerTensor) = blockperm(store(C))
 blockcomb(C::CombinerTensor) = blockcomb(store(C))
 
-Base.conj(T::CombinerTensor; always_copy = false) = T
+conj(T::CombinerTensor; always_copy = false) = T
 
 function contraction_output(::TensorT1,
                             ::TensorT2,
@@ -117,16 +117,12 @@ function contract!!(R::Tensor{<:Number,NR},
   return contract!!(R,labelsR,T2,labelsT2,T1,labelsT1)
 end
 
-function Base.show(io::IO,
-                   mime::MIME"text/plain",
-                   S::Combiner)
+function show(io::IO, mime::MIME"text/plain", S::Combiner)
   println(io, "Permutation of blocks: ", S.perm)
   println(io, "Combination of blocks: ", S.comb)
 end
 
-function Base.show(io::IO,
-                   mime::MIME"text/plain",
-                   T::CombinerTensor)
+function show(io::IO, mime::MIME"text/plain", T::CombinerTensor)
   summary(io, T)
   println(io)
   show(io, mime, store(T))

src/contraction_logic.jl

@@ -41,33 +41,62 @@ function _contract_inds!(Ris,
                          T2is,
                          T2labels::Labels{N2},
                          Rlabels::Labels{NR}) where {N1,N2,NR}
-  ncont = 0
-  for i in T1labels
-    i < 0 && (ncont += 1)
-  end
-  IndT = promote_type(eltype(T1is), eltype(T2is))
-  u = 1
-  # TODO: use Rlabels, don't assume ncon convention
-  for i1 ∈ 1:N1
-    if T1labels[i1] > 0
-      Ris[u] = T1is[i1]
-      u += 1 
-    else
-      # This is to check that T1is and T2is
-      # can contract
-      i2 = findfirst(==(T1labels[i1]),T2labels)
-      dir(T1is[i1]) == -dir(T2is[i2]) || error("Attempting to contract index:\n\n$(T1is[i1])\nwith index:\n\n$(T2is[i2])\nIndices must have opposite directions to contract.")
+  for n in 1:NR
+    Rlabel = @inbounds Rlabels[n]
+    found = false
+    for n1 in 1:N1
+      if Rlabel == @inbounds T1labels[n1]
+        @inbounds Ris[n] = @inbounds T1is[n1]
+        found = true
+        break
+      end
     end
-  end
-  for i2 ∈ 1:N2
-    if T2labels[i2] > 0
-      Ris[u] = T2is[i2]
-      u += 1 
+    if !found
+      for n2 in 1:N2
+        if Rlabel == @inbounds T2labels[n2]
+          @inbounds Ris[n] = @inbounds T2is[n2]
+          break
+        end
+      end
     end
   end
   return nothing
 end
 
+# Old version that doesn't take into account Rlabels
+#function _contract_inds!(Ris,
+#                         T1is,
+#                         T1labels::Labels{N1},
+#                         T2is,
+#                         T2labels::Labels{N2},
+#                         Rlabels::Labels{NR}) where {N1,N2,NR}
+#  ncont = 0
+#  for i in T1labels
+#    i < 0 && (ncont += 1)
+#  end
+#  IndT = promote_type(eltype(T1is), eltype(T2is))
+#  u = 1
+#  # TODO: use Rlabels, don't assume ncon convention
+#  for i1 ∈ 1:N1
+#    if T1labels[i1] > 0
+#      Ris[u] = T1is[i1]
+#      u += 1 
+#    else
+#      # This is to check that T1is and T2is
+#      # can contract
+#      i2 = findfirst(==(T1labels[i1]),T2labels)
+#      dir(T1is[i1]) == -dir(T2is[i2]) || error("Attempting to contract index:\n\n$(T1is[i1])\nwith index:\n\n$(T2is[i2])\nIndices must have opposite directions to contract.")
+#    end
+#  end
+#  for i2 ∈ 1:N2
+#    if T2labels[i2] > 0
+#      Ris[u] = T2is[i2]
+#      u += 1 
+#    end
+#  end
+#  return nothing
+#end
+
 function contract_inds(T1is,
                        T1labels::Labels{N1},
                        T2is,