From 669fce1de2750edfa985c024880dd55ecd94b9dc Mon Sep 17 00:00:00 2001
From: Daniel Seemaier <daniel+github@seemaier.de>
Date: Thu, 18 Apr 2024 17:18:09 +0200
Subject: [PATCH] fix: parallel bucket initialization

---
 kaminpar-shm/datastructures/csr_graph.h | 40 +++++++++++++++++++++++--
 kaminpar-shm/graphutils/permutator.cc   | 20 ++++++++++---
 2 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/kaminpar-shm/datastructures/csr_graph.h b/kaminpar-shm/datastructures/csr_graph.h
index 25ad68c0..c06997c2 100644
--- a/kaminpar-shm/datastructures/csr_graph.h
+++ b/kaminpar-shm/datastructures/csr_graph.h
@@ -451,10 +451,46 @@ class AbstractCSRGraph : public AbstractGraph {
   void init_degree_buckets() {
     KASSERT(std::all_of(_buckets.begin(), _buckets.end(), [](const auto n) { return n == 0; }));
 
+    constexpr std::size_t kNumBuckets = kNumberOfDegreeBuckets<NodeID> + 1;
+
     if (_sorted) {
-      for (const NodeID u : nodes()) {
-        ++_buckets[degree_bucket(degree(u)) + 1];
+      tbb::enumerable_thread_specific<std::array<NodeID, kNumBuckets>> buckets_ets([&] {
+        return std::array<NodeID, kNumBuckets>{};
+      });
+
+      tbb::parallel_for(
+          tbb::blocked_range<NodeID>(0, n()),
+          [&](const tbb::blocked_range<NodeID> r) {
+            auto &buckets = buckets_ets.local();
+            for (NodeID u = r.begin(); u != r.end(); ++u) {
+              ++buckets[degree_bucket(degree(u)) + 1];
+            }
+          }
+      );
+
+      std::fill(_buckets.begin(), _buckets.end(), 0);
+      for (auto &local_buckets : buckets_ets) {
+        for (std::size_t i = 0; i < kNumBuckets; ++i) {
+          _buckets[i] += local_buckets[i];
+        }
       }
+
+      KASSERT(
+          [&] {
+            std::vector<NodeID> buckets2(_buckets.size());
+            for (const NodeID u : nodes()) {
+              ++buckets2[degree_bucket(degree(u)) + 1];
+            }
+            for (std::size_t i = 0; i < _buckets.size(); ++i) {
+              if (_buckets[i] != buckets2[i]) {
+                return false;
+              }
+            }
+            return true;
+          }(),
+          "",
+          assert::heavy
+      );
       auto last_nonempty_bucket =
           std::find_if(_buckets.rbegin(), _buckets.rend(), [](const auto n) { return n > 0; });
       _number_of_buckets = std::distance(_buckets.begin(), (last_nonempty_bucket + 1).base());
diff --git a/kaminpar-shm/graphutils/permutator.cc b/kaminpar-shm/graphutils/permutator.cc
index ec67c383..734e4ceb 100644
--- a/kaminpar-shm/graphutils/permutator.cc
+++ b/kaminpar-shm/graphutils/permutator.cc
@@ -27,11 +27,13 @@ NodePermutations<StaticArray> rearrange_graph(
     StaticArray<EdgeWeight> &edge_weights
 ) {
   START_HEAP_PROFILER("Temporal nodes and edges allocation");
-  START_TIMER("Allocation");
+  START_TIMER("Allocation (noinit)");
   RECORD("tmp_nodes") StaticArray<EdgeID> tmp_nodes(nodes.size(), static_array::noinit);
   RECORD("tmp_edges") StaticArray<NodeID> tmp_edges(edges.size(), static_array::noinit);
-  RECORD("tmp_node_weights") StaticArray<NodeWeight> tmp_node_weights(node_weights.size(), static_array::noinit);
-  RECORD("tmp_edge_weights") StaticArray<EdgeWeight> tmp_edge_weights(edge_weights.size(), static_array::noinit);
+  RECORD("tmp_node_weights")
+  StaticArray<NodeWeight> tmp_node_weights(node_weights.size(), static_array::noinit);
+  RECORD("tmp_edge_weights")
+  StaticArray<EdgeWeight> tmp_edge_weights(edge_weights.size(), static_array::noinit);
   STOP_TIMER();
   STOP_HEAP_PROFILER();
 
@@ -39,8 +41,10 @@ NodePermutations<StaticArray> rearrange_graph(
   // the graph data structure this way, we can just cut them off without doing
   // further work
   START_HEAP_PROFILER("Rearrange input graph");
-  START_TIMER("Rearrange input graph");
+  START_TIMER("Sort nodes by degree bucket");
   NodePermutations<StaticArray> permutations = sort_by_degree_buckets<>(nodes);
+  STOP_TIMER();
+  START_TIMER("Rearrange input graph");
   build_permuted_graph(
       nodes,
       edges,
@@ -59,6 +63,13 @@ NodePermutations<StaticArray> rearrange_graph(
   STOP_TIMER();
   STOP_HEAP_PROFILER();
 
+  START_TIMER("Deallocation");
+  tmp_nodes.free();
+  tmp_edges.free();
+  tmp_node_weights.free();
+  tmp_edge_weights.free();
+  STOP_TIMER();
+
   return permutations;
 }
 
@@ -105,6 +116,7 @@ Graph rearrange_by_degree_buckets(CSRGraph &old_graph) {
       std::move(nodes), std::move(edges), std::move(node_weights), std::move(edge_weights), true
   ));
   new_graph.set_permutation(std::move(node_permutations.old_to_new));
+
   return new_graph;
 }