From d8e9cbfe278898d3a302f5f7d62dba946eafd573 Mon Sep 17 00:00:00 2001
From: Daniel Seemaier <daniel+github@seemaier.de>
Date: Thu, 26 Sep 2024 19:08:54 +0200
Subject: [PATCH] fix(shm): sequential allocs in the greedy balancer

---
 kaminpar-common/datastructures/binary_heap.h  | 26 ++++++++++++-------
 kaminpar-common/datastructures/marker.h       | 21 +++++++++++----
 .../refinement/balancer/greedy_balancer.cc    |  4 +--
 .../refinement/balancer/greedy_balancer.h     |  4 +--
 4 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/kaminpar-common/datastructures/binary_heap.h b/kaminpar-common/datastructures/binary_heap.h
index 0dc01241..1785ca59 100644
--- a/kaminpar-common/datastructures/binary_heap.h
+++ b/kaminpar-common/datastructures/binary_heap.h
@@ -460,7 +460,12 @@ template <typename Key> using BinaryMaxHeap = BinaryHeap<Key, binary_heap::max_h
 
 template <typename Key> using BinaryMinHeap = BinaryHeap<Key, binary_heap::min_heap_comparator>;
 
-template <typename ID, typename Key, template <typename> typename Comparator>
+template <
+    typename ID,
+    typename Key,
+    template <typename>
+    typename Comparator,
+    template <typename> typename Container = std::vector>
 class DynamicBinaryForest {
   static constexpr std::size_t kTreeArity = 4;
   static constexpr ID kInvalidID = std::numeric_limits<ID>::max();
@@ -645,18 +650,21 @@ class DynamicBinaryForest {
     std::swap(_heaps[heap][a], _heaps[heap][b]);
   }
 
-  std::vector<std::size_t> _id_pos;
+  Container<std::size_t> _id_pos;
   std::vector<std::vector<HeapElement>> _heaps;
   Comparator<Key> _comparator{};
 };
 
-template <typename ID, typename Key>
-using DynamicBinaryMaxForest = DynamicBinaryForest<ID, Key, binary_heap::max_heap_comparator>;
+template <typename ID, typename Key, template <typename> typename Container = std::vector>
+using DynamicBinaryMaxForest =
+    DynamicBinaryForest<ID, Key, binary_heap::max_heap_comparator, Container>;
 
-template <typename ID, typename Key>
-using DynamicBinaryMinForest = DynamicBinaryForest<ID, Key, binary_heap::min_heap_comparator>;
+template <typename ID, typename Key, template <typename> typename Container = std::vector>
+using DynamicBinaryMinForest =
+    DynamicBinaryForest<ID, Key, binary_heap::min_heap_comparator, Container>;
 
-template <typename ID, typename Key> class DynamicBinaryMinMaxForest {
+template <typename ID, typename Key, template <typename> typename Container = std::vector>
+class DynamicBinaryMinMaxForest {
 public:
   DynamicBinaryMinMaxForest() {}
 
@@ -763,8 +771,8 @@ template <typename ID, typename Key> class DynamicBinaryMinMaxForest {
   }
 
 private:
-  DynamicBinaryMaxForest<ID, Key> _max_forest;
-  DynamicBinaryMinForest<ID, Key> _min_forest;
+  DynamicBinaryMaxForest<ID, Key, Container> _max_forest;
+  DynamicBinaryMinForest<ID, Key, Container> _min_forest;
 };
 
 //! Dynamic binary heap, not addressable
diff --git a/kaminpar-common/datastructures/marker.h b/kaminpar-common/datastructures/marker.h
index 20c986ec..325fa6c7 100644
--- a/kaminpar-common/datastructures/marker.h
+++ b/kaminpar-common/datastructures/marker.h
@@ -12,11 +12,17 @@
 #include <limits>
 #include <vector>
 
+#include <tbb/parallel_for.h>
+
 #include "kaminpar-common/assert.h"
 #include "kaminpar-common/heap_profiler.h"
 
 namespace kaminpar {
-template <std::size_t kNumConcurrentMarkers = 1, typename Value = std::size_t> class Marker {
+template <
+    std::size_t kNumConcurrentMarkers = 1,
+    typename Value = std::size_t,
+    template <typename> typename Container = std::vector>
+class Marker {
 public:
   explicit Marker() : _marker_id(0), _first_unmarked_element{0} {
     RECORD_DATA_STRUCT(0, _struct);
@@ -80,9 +86,14 @@ template <std::size_t kNumConcurrentMarkers = 1, typename Value = std::size_t> c
 
     if ((_marker_id | ((1u << kNumConcurrentMarkers) - 1u)) == std::numeric_limits<Value>::max()) {
       _marker_id = 0;
-      const auto capacity = _data.size();
-      _data.clear();
-      _data.resize(capacity, 0);
+
+      if constexpr (std::is_same_v<Container<Value>, std::vector<Value>>) {
+        const auto capacity = _data.size();
+        _data.clear();
+        _data.resize(capacity, 0);
+      } else {
+        tbb::parallel_for<std::size_t>(0, _data.size(), [&](const std::size_t i) { _data[i] = 0; });
+      }
     }
   }
 
@@ -96,7 +107,7 @@ template <std::size_t kNumConcurrentMarkers = 1, typename Value = std::size_t> c
   }
 
 private:
-  std::vector<Value> _data;
+  Container<Value> _data;
   Value _marker_id;
   std::array<std::size_t, kNumConcurrentMarkers> _first_unmarked_element;
 
diff --git a/kaminpar-shm/refinement/balancer/greedy_balancer.cc b/kaminpar-shm/refinement/balancer/greedy_balancer.cc
index 16675615..fa755b22 100644
--- a/kaminpar-shm/refinement/balancer/greedy_balancer.cc
+++ b/kaminpar-shm/refinement/balancer/greedy_balancer.cc
@@ -438,10 +438,10 @@ template <typename Graph> class GreedyBalancerImpl {
   PartitionedGraph *_p_graph;
   const Graph *_graph;
 
-  DynamicBinaryMinMaxForest<NodeID, double> _pq;
+  DynamicBinaryMinMaxForest<NodeID, double, StaticArray> _pq;
   mutable tbb::enumerable_thread_specific<RatingMap<EdgeWeight, NodeID>> _rating_map;
   tbb::enumerable_thread_specific<std::vector<BlockID>> _feasible_target_blocks;
-  Marker<> _marker;
+  Marker<1, std::size_t, StaticArray> _marker;
   std::vector<BlockWeight> _pq_weight;
 
   Statistics _stats;
diff --git a/kaminpar-shm/refinement/balancer/greedy_balancer.h b/kaminpar-shm/refinement/balancer/greedy_balancer.h
index 7b26aa69..b199385f 100644
--- a/kaminpar-shm/refinement/balancer/greedy_balancer.h
+++ b/kaminpar-shm/refinement/balancer/greedy_balancer.h
@@ -22,10 +22,10 @@ namespace kaminpar::shm {
 template <typename Graph> class GreedyBalancerImpl;
 
 struct GreedyBalancerMemoryContext {
-  DynamicBinaryMinMaxForest<NodeID, double> pq;
+  DynamicBinaryMinMaxForest<NodeID, double, StaticArray> pq;
   tbb::enumerable_thread_specific<RatingMap<EdgeWeight, NodeID>> rating_map;
   tbb::enumerable_thread_specific<std::vector<BlockID>> feasible_target_blocks;
-  Marker<> marker;
+  Marker<1, std::size_t, StaticArray> marker;
   std::vector<BlockWeight> pq_weight;
   NormalSparseGainCache<Graph> *gain_cache = nullptr;
 };