feed num nodes to the rebalancer

kahypar · Jul 14, 2023 · 66d5a91 · 66d5a91
1 parent 27a4632
commit 66d5a91
Show file tree

Hide file tree

Showing 8 changed files with 30 additions and 33 deletions.
diff --git a/mt-kahypar/partition/coarsening/uncoarsener_base.h b/mt-kahypar/partition/coarsening/uncoarsener_base.h
@@ -128,7 +128,7 @@ class UncoarsenerBase {
       _context.refinement.flows.algorithm,
       _hg.initialNumNodes(), _hg.initialNumEdges(), _context, _gain_cache);
     _rebalancer = RebalancerFactory::getInstance().createObject(
-      _context.refinement.rebalancer, _context, _gain_cache);
+      _context.refinement.rebalancer, _hg.initialNumNodes(), _context, _gain_cache);
     // JET and FMrequires acces to the rebalancer
     _fm = FMFactory::getInstance().createObject(
       _context.refinement.fm.algorithm,

diff --git a/mt-kahypar/partition/factories.h b/mt-kahypar/partition/factories.h
@@ -149,7 +149,7 @@ using FlowSchedulerDispatcher = kahypar::meta::StaticMultiDispatchFactory<
                                   IRefiner,
                                   kahypar::meta::Typelist<TypeTraitsList, GainTypes>>;
 
-using RebalancerFactory = kahypar::meta::Factory<RebalancingAlgorithm, IRebalancer* (*)(const Context&, gain_cache_t)>;
+using RebalancerFactory = kahypar::meta::Factory<RebalancingAlgorithm, IRebalancer* (*)(HypernodeID, const Context&, gain_cache_t)>;
 
 using RebalancerDispatcher = kahypar::meta::StaticMultiDispatchFactory<
                               Rebalancer,

diff --git a/mt-kahypar/partition/refinement/fm/multitry_kway_fm.cpp b/mt-kahypar/partition/refinement/fm/multitry_kway_fm.cpp
@@ -32,7 +32,7 @@
 #include "mt-kahypar/datastructures/streaming_vector.h"
 #include "mt-kahypar/definitions.h"
 #include "mt-kahypar/utils/utilities.h"
-#include "mt-kahypar/partition/factories.h"
+#include "mt-kahypar/partition/factories.h"   // TODO removing this could make compilation a lot faster
 #include "mt-kahypar/partition/metrics.h"
 #include "mt-kahypar/partition/refinement/gains/gain_definitions.h"
 #include "mt-kahypar/utils/memory_tree.h"

diff --git a/mt-kahypar/partition/refinement/rebalancing/jet_rebalancer.h b/mt-kahypar/partition/refinement/rebalancing/jet_rebalancer.h
@@ -58,8 +58,7 @@ class JetRebalancer final : public IRebalancer {
 public:
   // TODO: add repairEmptyBlocks functionality?
 
-  explicit JetRebalancer(const Context& context,
-                         GainCache& gain_cache) :
+  explicit JetRebalancer(HypernodeID , const Context& context, GainCache& gain_cache) :
     _context(context),
     _max_part_weights(nullptr),
     _gain_cache(gain_cache),
@@ -81,9 +80,8 @@ class JetRebalancer final : public IRebalancer {
       }
     }
 
-  explicit JetRebalancer(const Context& context,
-                         gain_cache_t gain_cache) :
-    JetRebalancer(context, GainCachePtr::cast<GainCache>(gain_cache)) {}
+  explicit JetRebalancer(HypernodeID num_nodes, const Context& context, gain_cache_t gain_cache) :
+    JetRebalancer(num_nodes, context, GainCachePtr::cast<GainCache>(gain_cache)) {}
 
   JetRebalancer(const JetRebalancer&) = delete;
   JetRebalancer(JetRebalancer&&) = delete;

diff --git a/mt-kahypar/partition/refinement/rebalancing/rebalancer.h b/mt-kahypar/partition/refinement/rebalancing/rebalancer.h
@@ -76,11 +76,11 @@ class Rebalancer final : public IRebalancer {
     _gain(context),
     _part_weights(_context.partition.k) { }
 
-  explicit Rebalancer(const Context& context, GainCache&) :
+  explicit Rebalancer(HypernodeID , const Context& context, GainCache&) :
     Rebalancer(context) { }
 
-  explicit Rebalancer(const Context& context, gain_cache_t gain_cache) :
-    Rebalancer(context, GainCachePtr::cast<GainCache>(gain_cache)) {}
+  explicit Rebalancer(HypernodeID num_nodes, const Context& context, gain_cache_t gain_cache) :
+    Rebalancer(num_nodes, context, GainCachePtr::cast<GainCache>(gain_cache)) {}
 
   Rebalancer(const Rebalancer&) = delete;
   Rebalancer(Rebalancer&&) = delete;

diff --git a/mt-kahypar/partition/refinement/rebalancing/rebalancer_v2.cpp b/mt-kahypar/partition/refinement/rebalancing/rebalancer_v2.cpp
@@ -463,27 +463,24 @@ namespace impl {
   }
 
 
-
-
 template <typename TypeTraits, typename GainTypes>
-RebalancerV2<TypeTraits, GainTypes>::RebalancerV2(const Context& context,
-                      GainCache& gain_cache) :
+RebalancerV2<TypeTraits, GainTypes>::RebalancerV2(
+        HypernodeID num_nodes, const Context& context, GainCache& gain_cache) :
         _context(context),
         _max_part_weights(nullptr),
         _gain_cache(gain_cache),
         _current_k(_context.partition.k),
-        _gain(context)
-  // TODO add these
-  //target_part(num_nodes, kInvalidPartition),
-  //pq_handles(num_nodes, invalid_position),
-  //pq_id(num_nodes, -1),
-  //node_state(num_nodes) { }
-{ }
+        _gain(context),
+        _moves(num_nodes),
+        _target_part(num_nodes, kInvalidPartition),
+        _pq_handles(num_nodes, invalid_position),
+        _pq_id(num_nodes, -1),
+        _node_state(num_nodes) { }
 
 template <typename TypeTraits, typename GainTypes>
-RebalancerV2<TypeTraits, GainTypes>::RebalancerV2(const Context& context,
-                      gain_cache_t gain_cache) :
-        RebalancerV2(context, GainCachePtr::cast<GainCache>(gain_cache)) { }
+RebalancerV2<TypeTraits, GainTypes>::RebalancerV2(
+        HypernodeID num_nodes, const Context& context, gain_cache_t gain_cache) :
+        RebalancerV2(num_nodes, context, GainCachePtr::cast<GainCache>(gain_cache)) { }
 
 
 template <typename TypeTraits, typename GainTypes>

diff --git a/mt-kahypar/partition/refinement/rebalancing/rebalancer_v2.h b/mt-kahypar/partition/refinement/rebalancing/rebalancer_v2.h
@@ -81,11 +81,13 @@ class RebalancerV2 final : public IRebalancer {
 
 public:
 
-  explicit RebalancerV2(const Context& context,
-                         GainCache& gain_cache);
+  explicit RebalancerV2(HypernodeID num_nodes,
+                        const Context& context,
+                        GainCache& gain_cache);
 
-  explicit RebalancerV2(const Context& context,
-                         gain_cache_t gain_cache);
+  explicit RebalancerV2(HypernodeID num_nodes,
+                        const Context& context,
+                        gain_cache_t gain_cache);
 
 private:
   bool refineImpl(mt_kahypar_partitioned_hypergraph_t& hypergraph,

diff --git a/mt-kahypar/partition/registries/register_refinement_algorithms.cpp b/mt-kahypar/partition/registries/register_refinement_algorithms.cpp
@@ -121,18 +121,18 @@
 #define REGISTER_DISPATCHED_REBALANCER(id, dispatcher, ...)                                            \
   static kahypar::meta::Registrar<RebalancerFactory> register_ ## dispatcher(                          \
     id,                                                                                                \
-    [](const Context& context, gain_cache_t gain_cache) {                                              \
+    [](HypernodeID num_hypernodes, const Context& context, gain_cache_t gain_cache) {                                              \
     return dispatcher::create(                                                                         \
-      std::forward_as_tuple(context, gain_cache),                                                      \
+      std::forward_as_tuple(num_hypernodes, context, gain_cache),                                                      \
       __VA_ARGS__                                                                                      \
       );                                                                                               \
   })
 
 #define REGISTER_REBALANCER(id, refiner, t)                                                      \
   static kahypar::meta::Registrar<RebalancerFactory> JOIN(register_ ## refiner, t)(              \
     id,                                                                                          \
-    [](const Context& context, gain_cache_t gain_cache) -> IRebalancer* {                           \
-    return new refiner(context, gain_cache);                                                     \
+    [](HypernodeID num_hypernodes, const Context& context, gain_cache_t gain_cache) -> IRebalancer* {                           \
+    return new refiner(num_hypernodes, context, gain_cache);                                                     \
   })
 
 #define REGISTER_DISPATCHED_FLOW_REFINER(id, dispatcher, ...)                                          \