Fix rebalance init

kahypar · Dec 20, 2024 · df98874 · df98874
1 parent a932aea
commit df98874
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/mt-kahypar/dynamic/strategies/localFM_factor.h b/mt-kahypar/dynamic/strategies/localFM_factor.h
@@ -59,6 +59,7 @@ namespace mt_kahypar::dyn {
           //TODO: is second reset after rebalancing necessary?
           GainCachePtr::resetGainCache(_gain_cache);
 
+          //TODO does this need to be done prior to rebalancing?
           _fm->initialize(partitioned_hypergraph);
 
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
@@ -140,7 +141,7 @@ namespace mt_kahypar::dyn {
           if (skipped_changes >= step_size) {
             skipped_changes = 0;
             step_size *= 2;
-            local_fm(hypergraph, context);
+            local_fm(hypergraph, context, nodes_to_partition);
             nodes_to_partition = parallel::scalable_vector<HypernodeID>(changes_size);
           } else {
             skipped_changes++;

diff --git a/mt-kahypar/dynamic/strategies/rebalance.h b/mt-kahypar/dynamic/strategies/rebalance.h
@@ -13,6 +13,7 @@ namespace mt_kahypar::dyn {
         std::optional<ds::PartitionedHypergraph<ds::StaticHypergraph>> partitioned_hypergraph_s;
         gain_cache_t _gain_cache;
         std::unique_ptr<IRebalancer> _rebalancer;
+        std::unique_ptr<IRefiner> _fm;
 
         void repartition(ds::StaticHypergraph& hypergraph_s, Context& context) {
           context.dynamic.repartition_count++;
@@ -26,10 +27,19 @@ namespace mt_kahypar::dyn {
           _gain_cache = GainCachePtr::constructGainCache(context);
           _rebalancer = RebalancerFactory::getInstance().createObject(
                   context.refinement.rebalancer, hypergraph_s.initialNumNodes(), context, _gain_cache);
+
+          context.refinement.fm.algorithm = FMAlgorithm::kway_fm;
+          context.refinement.fm.multitry_rounds = context.dynamic.multitry_localFM;
+
+          _fm = FMFactory::getInstance().createObject(
+                  context.refinement.fm.algorithm,
+                  hypergraph_s.initialNumNodes(), hypergraph_s.initialNumEdges(), context, _gain_cache, *_rebalancer);
         }
 
         //use rebalancer to rebalance partitioned_hypergraph_s
-        void rebalance(ds::StaticHypergraph& hypergraph, Context& context) {
+        void rebalance(Context& context) {
+
+          std::cout << "Rebalancing" << std::endl;
 
           GainCachePtr::resetGainCache(_gain_cache);
 
@@ -38,6 +48,8 @@ namespace mt_kahypar::dyn {
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
                                   mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
 
+          _fm->initialize(partitioned_hypergraph);
+
           _rebalancer->refineAndOutputMoves(partitioned_hypergraph, {}, moves_by_part, best_Metrics, std::numeric_limits<double>::max());
         }
 
@@ -93,7 +105,7 @@ namespace mt_kahypar::dyn {
           }
 
           if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
-            rebalance(hypergraph, context);
+            rebalance(context);
           }
 
           ASSERT(metrics::isBalanced(*partitioned_hypergraph_s, context));