fix: flow non-determinism with steiner_tree metric (#221)

Fixes a rare case where the deterministic_quality preset could produce non-deterministic results when using the steiner_tree metric

Author: N-Maas

mt-kahypar/partition/refinement/flows/deterministic/deterministic_flow_refinement_scheduler.cpp

@@ -74,43 +74,78 @@ bool DeterministicFlowRefinementScheduler<GraphAndGainTypes>::refineImpl(mt_kahy
   HyperedgeWeight min_improvement = _context.refinement.flows.min_relative_improvement_per_round * current_metrics.quality;
   std::atomic<HyperedgeWeight> round_delta(std::numeric_limits<HyperedgeWeight>::max());
   tbb::concurrent_queue<BlockPair> scheduled_blocks;
+  vec<std::pair<BlockPair, MoveSequence>> resulting_moves;
+  std::atomic<size_t> result_id;
   for (size_t round = 0; round_delta >= min_improvement; ++round) {
     size_t num_scheduled_blocks = _schedule.getNextMatching(scheduled_blocks);
     if (num_scheduled_blocks == 0) break;
 
     round_delta = 0;
     while (num_scheduled_blocks > 0) {
       DBG << "Next matching:" << num_scheduled_blocks << "blocks";
-      _new_cut_hes.clear_sequential();
+      resulting_moves.assign(num_scheduled_blocks, {});
+      resulting_moves.resize(num_scheduled_blocks);
+      result_id = 0;
 
+      // Important:
+      // The next step runs the actual flow computations, but doesnt't apply the results yet.
+      // Applying the moves immediately would cause a determinism bug with the steiner_tree
+      // objective! This is because with the steiner_tree metric, the gain for moving pins of
+      // a hyperedge is influenced by all blocks in the connectivity set (not only source and
+      // target block). Therefore, applying moves constitutes a race condition with all threads
+      // that currently construct a flow problem which includes a part of the same hyperedge -
+      // which can happen even if we only schedule matchings. We solve this by applying the
+      // moves afterwards in a separate step.
       tbb::parallel_for(UL(0), _context.refinement.flows.num_parallel_searches, [&](const size_t refiner_idx) {
         BlockPair blocks;
         while (scheduled_blocks.try_pop(blocks)) {
           timer.start_timer("region_growing", "Grow Region", true);
           const Subhypergraph sub_hg = _constructor.construct(blocks, _quotient_graph, phg, /*deterministic=*/true);
           timer.stop_timer("region_growing");
 
-          HyperedgeWeight improvement = 0;
+          MoveSequence moves;
+          size_t local_result_id = result_id.fetch_add(1, std::memory_order_relaxed);
           auto start = std::chrono::high_resolution_clock::now();
           if ( sub_hg.numNodes() > 0 ) {
             auto partitioned_hg = utils::partitioned_hg_const_cast(phg);
             if (_refiner[refiner_idx] == nullptr) {
               _refiner[refiner_idx] = constructFlowRefiner();
             }
             _refiner[refiner_idx]->initialize(partitioned_hg);
-            MoveSequence moves = _refiner[refiner_idx]->refine(partitioned_hg, sub_hg, start);
+            moves = _refiner[refiner_idx]->refine(partitioned_hg, sub_hg, start);
+          }
+          ASSERT(local_result_id < resulting_moves.size());
+          resulting_moves[local_result_id] = {blocks, std::move(moves)};
+        }
+      });
 
-            timer.start_timer("apply_moves", "Apply Moves", true);
-            improvement = applyMoves(blocks, moves, phg);
-            timer.stop_timer("apply_moves");
+      // apply the resulting moves
+      timer.start_timer("apply_moves", "Apply Moves");
+      auto apply_moves = [&](BlockPair blocks, const MoveSequence& moves) {
+        ASSERT(blocks.i != kInvalidPartition && blocks.j != kInvalidPartition);
+        HyperedgeWeight improvement = applyMoves(blocks, moves, phg);
 
-            _schedule.reportResults(blocks, improvement);
-          }
+        round_delta.fetch_add(improvement, std::memory_order_relaxed);
+        _schedule.reportResults(blocks, improvement);
+        _quotient_graph.edge(blocks).reportImprovement(improvement);
+      };
 
-          round_delta.fetch_add(improvement, std::memory_order_relaxed);
-          _quotient_graph.edge(blocks).reportImprovement(improvement);
+      _new_cut_hes.clear_sequential();
+      if constexpr (FlowNetworkConstruction::is_exact_model) {
+        tbb::parallel_for(UL(0), num_scheduled_blocks, [&](const size_t i) {
+          apply_moves(resulting_moves[i].first, resulting_moves[i].second);
+        });
+      } else {
+        // in case of the steiner_tree metric, we need to apply the moves sequentially
+        // to ensure the reported improvement for each quotient graph edge is deterministic
+        std::sort(resulting_moves.begin(), resulting_moves.end(), [](auto& lhs, auto& rhs) {
+          return std::tie(lhs.first.i, lhs.first.j) < std::tie(rhs.first.i, rhs.first.j);
+        });
+        for (const auto& pair: resulting_moves) {
+          apply_moves(pair.first, pair.second);
         }
-      });
+      }
+      timer.stop_timer("apply_moves");
 
       // adding HEs to the quotient graph must happen after all moves are applied,
       // since the result depends on the current connectivity set of the HE
@@ -200,7 +235,7 @@ bool changeNodePart(PartitionedHypergraph& phg,
 
 template<typename GraphAndGainTypes>
 HyperedgeWeight DeterministicFlowRefinementScheduler<GraphAndGainTypes>::applyMoves(const BlockPair& blocks,
-                                                                                    MoveSequence& sequence,
+                                                                                    const MoveSequence& sequence,
                                                                                     PartitionedHypergraph& phg) {
   const bool gain_cache_update = _context.forceGainCacheUpdates();
   HyperedgeWeight improvement = 0;

mt-kahypar/partition/refinement/flows/deterministic/deterministic_flow_refinement_scheduler.h

@@ -84,7 +84,7 @@ class DeterministicFlowRefinementScheduler final : public IRefiner {
 
   std::unique_ptr<IFlowRefiner> constructFlowRefiner();
 
-  HyperedgeWeight applyMoves(const BlockPair& blocks, MoveSequence& sequence, PartitionedHypergraph& phg);
+  HyperedgeWeight applyMoves(const BlockPair& blocks, const MoveSequence& sequence, PartitionedHypergraph& phg);
 
   void initializeImpl(mt_kahypar_partitioned_hypergraph_t& hypergraph);
 

mt-kahypar/partition/refinement/flows/flow_common.h

@@ -46,7 +46,7 @@ enum class MoveSequenceState : uint8_t {
 // apply the moves
 struct MoveSequence {
   vec<Move> moves;
-  Gain expected_improvement; // >= 0
+  Gain expected_improvement = 0; // >= 0
   MoveSequenceState state = MoveSequenceState::IN_PROGRESS;
 };
 

mt-kahypar/partition/refinement/gains/steiner_tree/steiner_tree_flow_network_construction.cpp

@@ -121,7 +121,7 @@ FlowNetworkEdgeParameters SteinerTreeFlowNetworkConstruction::getParameters(cons
     const HyperedgeWeight distance_without_block_1 = target_graph.distanceWithoutBlock(connectivity_set, block_1);
     const HyperedgeWeight gain_0 = (current_distance - distance_without_block_0) * edge_weight;
     const HyperedgeWeight gain_1 = (current_distance - distance_without_block_1) * edge_weight;
-    parameters.capacity = capacity_for_cut_edge(context.refinement.flows.steiner_tree_policy, gain_0, gain_1);
+    parameters.capacity = std::max(0, capacity_for_cut_edge(context.refinement.flows.steiner_tree_policy, gain_0, gain_1));
   }
   return parameters;
 }