Smarter Imbalance Checks

kahypar · Dec 19, 2024 · d1172de · d1172de
1 parent 0ad9e9d
commit d1172de
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git a/mt-kahypar/dynamic/strategies/connectivity.h b/mt-kahypar/dynamic/strategies/connectivity.h
@@ -90,8 +90,7 @@ namespace mt_kahypar::dyn {
 
 
           //check if imbalance is still within bounds else repartition
-          if (partitioned_hypergraph_s->partWeight(std::get<1>(*max_connectivity)) /
-              static_cast<double>(context.partition.perfect_balance_part_weights[std::get<1>(*max_connectivity)]) - 1.0 > context.partition.epsilon ) {
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
             repartition(hypergraph, context);
           }
           history.push_back(partition_result);

diff --git a/mt-kahypar/dynamic/strategies/localFM.h b/mt-kahypar/dynamic/strategies/localFM.h
@@ -46,7 +46,7 @@ namespace mt_kahypar::dyn {
 
           mt_kahypar_partitioned_hypergraph_t  partitioned_hypergraph = utils::partitioned_hg_cast(*partitioned_hypergraph_s);
 
-          if (mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context) > context.partition.epsilon) {
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
             // use rebalancer to rebalance partitioned_hypergraph_s
             parallel::scalable_vector<parallel::scalable_vector<Move>> moves_by_part;
             Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
@@ -127,10 +127,8 @@ namespace mt_kahypar::dyn {
 
           local_fm(hypergraph, context, local_fm_nodes);
 
-          //check if imbalance is still within bounds else repartition
-          if (mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context) > context.partition.epsilon + 0.0000001) {
-            repartition(hypergraph, context);
-          }
+          ASSERT(metrics::isBalanced(*partitioned_hypergraph_s, context));
+
           history.push_back(partition_result);
         }
 

diff --git a/mt-kahypar/dynamic/strategies/localFM_factor.h b/mt-kahypar/dynamic/strategies/localFM_factor.h
@@ -42,18 +42,30 @@ namespace mt_kahypar::dyn {
         }
 
         //use local_fm to refine partitioned_hypergraph_s
-        void local_fm(ds::StaticHypergraph& hypergraph, Context& context) {
+        void local_fm(ds::StaticHypergraph& hypergraph, Context& context, parallel::scalable_vector<HypernodeID> local_fm_nodes) {
 
           GainCachePtr::resetGainCache(_gain_cache);
 
           mt_kahypar_partitioned_hypergraph_t  partitioned_hypergraph = utils::partitioned_hg_cast(*partitioned_hypergraph_s);
 
+          if (!metrics::isBalanced(*partitioned_hypergraph_s, context)) {
+            // use rebalancer to rebalance partitioned_hypergraph_s
+            parallel::scalable_vector<parallel::scalable_vector<Move>> moves_by_part;
+            Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
+                                    mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
+            _rebalancer->refineAndOutputMoves(partitioned_hypergraph, {}, moves_by_part, best_Metrics, std::numeric_limits<double>::max());
+          }
+
+          //TODO: is second reset after rebalancing necessary?
+          GainCachePtr::resetGainCache(_gain_cache);
+
           _fm->initialize(partitioned_hypergraph);
 
           Metrics best_Metrics = {mt_kahypar::metrics::quality(*partitioned_hypergraph_s, Objective::km1),
                                   mt_kahypar::metrics::imbalance(*partitioned_hypergraph_s, context)};
 
-          _fm->refine(partitioned_hypergraph, nodes_to_partition, best_Metrics, std::numeric_limits<double>::max());
+          _fm->refine(partitioned_hypergraph, local_fm_nodes, best_Metrics, std::numeric_limits<double>::max());
+
         }
 
         PartitionID add_node_to_partitioned_hypergraph(ds::StaticHypergraph& hypergraph, Context& context, const HypernodeID& hn) {