potentially more efficient sorting

kahypar · Nov 17, 2023 · e7dfb74 · e7dfb74
1 parent 9428631
commit e7dfb74
Showing 1 changed file with 21 additions and 4 deletions.
diff --git a/mt-kahypar/datastructures/static_graph.cpp b/mt-kahypar/datastructures/static_graph.cpp
@@ -33,6 +33,8 @@
 #include "mt-kahypar/utils/timer.h"
 #include "mt-kahypar/utils/memory_tree.h"
 
+#include <algorithm>
+
 #include <tbb/parallel_reduce.h>
 #include <tbb/parallel_sort.h>
 
@@ -259,8 +261,9 @@ namespace mt_kahypar::ds {
       }
       incident_edges_map.reserve_for_estimated_number_of_insertions(max_degree);
 
+      parallel::scalable_vector<std::pair<size_t, size_t>> resulting_ranges;
+      timer.start_timer("stage_3_hd", "Contract Stage 3 (HD)");
       for ( const HypernodeID& coarse_node : high_degree_vertices ) {
-        timer.start_timer("stage_3_hd", "Contract Stage 3 (HD)");
         const size_t incident_edges_start = tmp_incident_edges_prefix_sum[coarse_node];
         const size_t incident_edges_end = tmp_incident_edges_prefix_sum[coarse_node + 1];
         const size_t size_of_range = incident_edges_end - incident_edges_start;
@@ -301,16 +304,30 @@ namespace mt_kahypar::ds {
           tmp_edges[i].invalidate();
         });
         node_sizes[coarse_node] = contracted_size;
-        timer.stop_timer("stage_3_hd");
+
+        resulting_ranges.emplace_back(incident_edges_start, incident_edges_pos.load());
 
         // sort
-        timer.start_timer("stage_3_sort", "Contract Stage 3 (sort)");
         tbb::parallel_sort(tmp_edges.begin() + incident_edges_start, tmp_edges.begin() + incident_edges_pos.load(),
           [](const TmpEdgeInformation& e1, const TmpEdgeInformation& e2) {
             return e1._target < e2._target;
           });
-        timer.stop_timer("stage_3_sort");
       }
+      timer.stop_timer("stage_3_hd");
+
+      timer.start_timer("stage_3_sort", "Contract Stage 3 (sort)");
+      tbb::parallel_for(UL(0), resulting_ranges.size(), [&](const size_t i) {
+        auto [start, end] = resulting_ranges[i];
+        auto comparator = [](const TmpEdgeInformation& e1, const TmpEdgeInformation& e2) {
+          return e1._target < e2._target;
+        };
+        if (end - start <= high_degree_threshold) {
+          std::sort(tmp_edges.begin() + start, tmp_edges.begin() + end, comparator);
+        } else {
+          tbb::parallel_sort(tmp_edges.begin() + start, tmp_edges.begin() + end, comparator);
+        }
+      });
+      timer.stop_timer("stage_3_sort");
     }
 
     // #################### STAGE 4 ####################