Better performance and type stability

Author: mtfishman

NDTensors/src/NDTensors.jl

@@ -72,8 +72,7 @@ include("blocksparse/contract.jl")
 include("blocksparse/contract_utilities.jl")
 include("blocksparse/contract_generic.jl")
 include("blocksparse/contract_sequential.jl")
-include("blocksparse/contract_folds.jl")
-include("blocksparse/contract_threads.jl")
+include("blocksparse/contract_threaded.jl")
 include("blocksparse/diagblocksparse.jl")
 include("blocksparse/similar.jl")
 include("blocksparse/combiner.jl")

NDTensors/src/blocksparse/contract_folds.jl

@@ -11,8 +11,6 @@ function contract_blockoffsets(
   indsR,
   labelsR,
 )
-  N1 = length(blocktype(boffs1))
-  N2 = length(blocktype(boffs2))
   NR = length(labelsR)
   ValNR = ValLength(labelsR)
   labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels(

NDTensors/src/blocksparse/contract_generic.jl

@@ -0,0 +1,78 @@
+using .Expose: expose
+function contract_blocks(
+  alg::Algorithm"threaded_threads",
+  boffs1,
+  boffs2,
+  labels1_to_labels2,
+  labels1_to_labelsR,
+  labels2_to_labelsR,
+  ValNR::Val{NR},
+) where {NR}
+  N1 = length(blocktype(boffs1))
+  N2 = length(blocktype(boffs2))
+  blocks1 = keys(boffs1)
+  blocks2 = keys(boffs2)
+  return if length(blocks1) > length(blocks2)
+    tasks = map(
+      Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads()))
+    ) do blocks1_partition
+      @spawn begin
+        block_contractions = Tuple{Block{N1},Block{N2},Block{NR}}[]
+        foreach(Iterators.product(blocks1_partition, blocks2)) do (block1, block2)
+          block_contraction = maybe_contract_blocks(
+            block1,
+            block2,
+            labels1_to_labels2,
+            labels1_to_labelsR,
+            labels2_to_labelsR,
+            ValNR,
+          )
+          if !isnothing(block_contraction)
+            push!(block_contractions, block_contraction)
+          end
+        end
+        return block_contractions
+      end
+    end
+    mapreduce(fetch, vcat, tasks)
+  else
+    tasks = map(
+      Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads()))
+    ) do blocks2_partition
+      @spawn begin
+        block_contractions = Tuple{Block{N1},Block{N2},Block{NR}}[]
+        foreach(Iterators.product(blocks1, blocks2_partition)) do (block1, block2)
+          block_contraction = maybe_contract_blocks(
+            block1,
+            block2,
+            labels1_to_labels2,
+            labels1_to_labelsR,
+            labels2_to_labelsR,
+            ValNR,
+          )
+          if !isnothing(block_contraction)
+            push!(block_contractions, block_contraction)
+          end
+        end
+        return block_contractions
+      end
+    end
+    mapreduce(fetch, vcat, tasks)
+  end
+end
+
+function contract!(
+  ::Algorithm"threaded_folds",
+  R::BlockSparseTensor,
+  labelsR,
+  tensor1::BlockSparseTensor,
+  labelstensor1,
+  tensor2::BlockSparseTensor,
+  labelstensor2,
+  contraction_plan,
+)
+  executor = ThreadedEx()
+  return contract!(
+    R, labelsR, tensor1, labelstensor1, tensor2, labelstensor2, contraction_plan, executor
+  )
+end

NDTensors/src/blocksparse/contract_threaded.jl

@@ -1,4 +1,3 @@
-using Compat
 using ITensors
 using Test
 using LinearAlgebra
@@ -8,7 +7,7 @@ if isone(Threads.nthreads())
 end
 
 @testset "Threading" begin
-  blas_num_threads = Compat.get_num_threads()
+  blas_num_threads = BLAS.get_num_threads()
   strided_num_threads = ITensors.NDTensors.Strided.get_num_threads()
 
   BLAS.set_num_threads(1)