Rewrite threaded blockspares to avoid threadid

Author: mtfishman

NDTensors/src/blocksparse/contract_generic.jl

@@ -18,40 +18,9 @@ function contract_blockoffsets(
   labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels(
     labels1, labels2, labelsR
   )
-
-  # Contraction plan element type
-  T = Tuple{Block{N1},Block{N2},Block{NR}}
-
-  # Thread-local collections of block contractions.
-  # Could use:
-  # ```julia
-  # FLoops.@reduce(contraction_plans = append!(T[], [(block1, block2, blockR)]))
-  # ```
-  # as a simpler alternative but it is slower.
-
-  contraction_plans = Vector{T}[T[] for _ in 1:nthreads()]
-
-  #
-  # Reserve some capacity
-  # In theory the maximum is length(boffs1) * length(boffs2)
-  # but in practice that is too much
-  #for contraction_plan in contraction_plans
-  #  sizehint!(contraction_plan, max(length(boffs1), length(boffs2)))
-  #end
-  #
-
-  contract_blocks!(
-    alg,
-    contraction_plans,
-    boffs1,
-    boffs2,
-    labels1_to_labels2,
-    labels1_to_labelsR,
-    labels2_to_labelsR,
-    ValNR,
+  contraction_plan = contract_blocks(
+    alg, boffs1, boffs2, labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR, ValNR
   )
-
-  contraction_plan = reduce(vcat, contraction_plans)
   blockoffsetsR = BlockOffsets{NR}()
   nnzR = 0
   for (_, _, blockR) in contraction_plan
@@ -60,7 +29,6 @@ function contract_blockoffsets(
       nnzR += blockdim(indsR, blockR)
     end
   end
-
   return blockoffsetsR, contraction_plan
 end
 

NDTensors/src/blocksparse/contract_threads.jl

@@ -1,7 +1,6 @@
 using .Expose: expose
-function contract_blocks!(
+function contract_blocks(
   alg::Algorithm"threaded_threads",
-  contraction_plans,
   boffs1,
   boffs2,
   labels1_to_labels2,
@@ -11,40 +10,45 @@ function contract_blocks!(
 )
   blocks1 = keys(boffs1)
   blocks2 = keys(boffs2)
-  if length(blocks1) > length(blocks2)
-    @sync for blocks1_partition in
-              Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads()))
-      @spawn for block1 in blocks1_partition
-        for block2 in blocks2
-          maybe_contract_blocks!(
-            contraction_plans[threadid()],
-            block1,
-            block2,
-            labels1_to_labels2,
-            labels1_to_labelsR,
-            labels2_to_labelsR,
-            ValNR,
-          )
-        end
+  return if length(blocks1) > length(blocks2)
+    tasks = map(
+      Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads()))
+    ) do blocks1_partition
+      @spawn begin
+        block_contractions =
+          map(Iterators.product(blocks1_partition, blocks2)) do (block1, block2)
+            maybe_contract_blocks(
+              block1,
+              block2,
+              labels1_to_labels2,
+              labels1_to_labelsR,
+              labels2_to_labelsR,
+              ValNR,
+            )
+          end
+        block_contractions = filter(!isnothing, block_contractions)
       end
     end
+    mapreduce(fetch, vcat, tasks)
   else
-    @sync for blocks2_partition in
-              Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads()))
-      @spawn for block2 in blocks2_partition
-        for block1 in blocks1
-          maybe_contract_blocks!(
-            contraction_plans[threadid()],
-            block1,
-            block2,
-            labels1_to_labels2,
-            labels1_to_labelsR,
-            labels2_to_labelsR,
-            ValNR,
-          )
-        end
+    tasks = map(
+      Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads()))
+    ) do blocks2_partition
+      @spawn begin
+        block_contractions =
+          map(Iterators.product(blocks1, blocks2_partition)) do (block1, block2)
+            maybe_contract_blocks(
+              block1,
+              block2,
+              labels1_to_labels2,
+              labels1_to_labelsR,
+              labels2_to_labelsR,
+              ValNR,
+            )
+          end
+        block_contractions = filter(!isnothing, block_contractions)
       end
     end
+    mapreduce(fetch, vcat, tasks)
   end
-  return nothing
 end

NDTensors/src/blocksparse/contract_utilities.jl

@@ -14,18 +14,12 @@ function compute_alpha(
   return one(ElR)
 end
 
-function maybe_contract_blocks!(
-  contraction_plan,
-  block1,
-  block2,
-  labels1_to_labels2,
-  labels1_to_labelsR,
-  labels2_to_labelsR,
-  ValNR,
+function maybe_contract_blocks(
+  block1, block2, labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR, ValNR
 )
   if are_blocks_contracted(block1, block2, labels1_to_labels2)
     blockR = contract_blocks(block1, labels1_to_labelsR, block2, labels2_to_labelsR, ValNR)
-    push!(contraction_plan, (block1, block2, blockR))
+    return block1, block2, blockR
   end
   return nothing
 end