Merge pull request #4860 from vgteam/parallel-payload-caching

Parallel payload caching

Author: adamnovak

deps/libbdsg

@@ -442,22 +442,37 @@ void cache_payloads(
 
     const handlegraph::HandleGraph* graph_ptr = (const handlegraph::HandleGraph*) &gbz.graph;
 
+    double total_zipcode_time = 0.0, total_decoder_time = 0.0;
+    std::atomic<uint64_t> node_count = 0;
     gbz.graph.for_each_handle([&](const handle_t& handle) {
         nid_t node_id = gbz.graph.get_id(handle);
-        ZipCode zipcode;
         pos_t pos = make_pos_t(node_id, false, 0);
-        zipcode.fill_in_zipcode_from_pos(distance_index, pos, true, graph_ptr);
+        ZipCode zipcode;
+        zipcode.fill_in_zipcode_from_pos(distance_index, pos, false, graph_ptr);
+        zipcode.fill_in_full_decoder();
+        if (++node_count % 10000 == 0 && progress) {
+            double telapsed = gbwt::readTimer() - start;
+            #pragma omp critical (cerr)
+            std::cerr << "  Cached " << node_count << " nodes in " << telapsed << "s" << std::endl;
+        }
+
         payload_t payload = zipcode.get_payload_from_zip();
         if (payload == MIPayload::NO_CODE && oversized_zipcodes != nullptr) {
             // The zipcode is too large for the payload field.
             // Add it to the oversized zipcode list.
-            zipcode.fill_in_full_decoder();
-            size_t offset = oversized_zipcodes->size();
-            oversized_zipcodes->emplace_back(zipcode);
+            size_t offset;
+            #pragma omp critical (cache_payloads_zipcodes)
+            {
+                offset = oversized_zipcodes->size();
+                oversized_zipcodes->emplace_back(zipcode);
+            }
             payload = { 0, offset };
         }
-        node_id_to_payload.emplace(node_id, payload);
-    });
+        #pragma omp critical (cache_payloads_map)
+        {
+            node_id_to_payload.emplace(node_id, payload);
+        }
+    }, true);
 
     if (progress) {
         double seconds = gbwt::readTimer() - start;
@@ -521,6 +536,15 @@ gbwtgraph::DefaultMinimizerIndex build_minimizer_index(
         // A zipcode only depends on the node id.
         vg::hash_map<nid_t, payload_t> node_id_to_payload;
         node_id_to_payload.reserve(gbz.graph.max_node_id() - gbz.graph.min_node_id());
+        // Preload the distance index right before we use it.
+        // find_frequent_kmers uses a lot of memory/IO scanning the whole graph
+        // and might evict the mmap'd index pages from the OS page cache,
+        // causing cache_payloads to page-fault on every node. So we preload
+        // after kmer counting to ensure the index is warm.
+        if (params.progress) {
+            std::cerr << "Preloading distance index";
+        }
+        distance_index->preload(true);
         cache_payloads(gbz, *distance_index, node_id_to_payload, oversized_zipcodes, params.progress);
 
         auto get_payload = [&](const pos_t& pos) -> const code_type* {

src/gbwtgraph_helper.cpp

@@ -814,10 +814,12 @@ static void populate_hub_labeling(SnarlDistanceIndex::TemporaryDistanceIndex& te
 /**
  * Determine if a snarl is regular or not.
  *
- * A regular snarl is a snarl that, while not simple, consists of only nodes or
+ * A regular snarl is a snarl that consists of only nodes or
  * chains connected to the start and end, without any connections between
  * multiple children, or any way to turn around. There may be an edge directly
  * across.
+ *
+ * A simple snarl is always regular.
  */
 static bool check_regularity(const SnarlDistanceIndex::TemporaryDistanceIndex& temp_index, const SnarlDistanceIndex::temp_record_ref_t& snarl_index, const SnarlDistanceIndex::TemporaryDistanceIndex::TemporarySnarlRecord& temp_snarl_record, const vector<SnarlDistanceIndex::temp_record_ref_t>& all_children, const HandleGraph* graph);
 
@@ -1661,11 +1663,11 @@ bool check_regularity(const SnarlDistanceIndex::TemporaryDistanceIndex& temp_ind
         return false;
     }
     if (temp_snarl_record.is_simple) {
-        // Simple snarls can't be regular because simple is more specific and useful.
+        // Simple snarls are always also regular.
 #ifdef debug_distance_indexing
-        std::cerr << "Snarl is not regular because it is simple." << std::endl;
+        std::cerr << "Snarl is regular because it is simple." << std::endl;
 #endif
-        return false;
+        return true;
     }
 
     // Get the snarl boundary nodes, facing out

src/snarl_distance_index.cpp

@@ -88,10 +88,11 @@ int main_minimizer(int argc, char** argv) {
     if (!config.distance_name.empty()) {
         // new distance index
         if (config.progress) {
-            logger.info() << "Loading SnarlDistanceIndex from " << config.distance_name << std::endl;
+            logger.info() << "Opening SnarlDistanceIndex at " << config.distance_name << std::endl;
         }
         distance_index = vg::io::VPKG::load_one<SnarlDistanceIndex>(config.distance_name);
-        distance_index->preload(true);
+        // Note that we don't fault in the index until we're actually about to
+        // use it, or it might get paged out again.
     }
 
     ZipCodeCollection oversized_zipcodes;

src/subcommand/minimizer_main.cpp

@@ -117,6 +117,10 @@ using namespace std;
         bool chain_is_reversed = distance_index.is_reversed_in_parent(
                                                 distance_index.get_node_net_handle(n1->id()));
 
+        // Node 4 is in snarl 3 to 6 which should be regular.
+        // The zip codes are going to encode this so it had better be true.
+        REQUIRE(distance_index.is_regular_snarl(distance_index.get_parent(distance_index.get_parent(distance_index.get_node_net_handle(n4->id())))));
+
         SECTION ("zip code for node on top-level chain") {
             ZipCode zipcode;
             zipcode.fill_in_zipcode_from_pos(distance_index, make_pos_t(n1->id(), 0, false));

src/unittest/zip_code.cpp

@@ -1,3 +1,5 @@
+#include "crash.hpp"
+
 #include "zip_code.hpp"
 
 //#define DEBUG_ZIPCODE
@@ -16,10 +18,11 @@ void ZipCode::fill_in_zipcode_from_pos(const SnarlDistanceIndex& distance_index,
     //Put all ancestors of the node in a vector, starting from the node, and not including the root
     while (!distance_index.is_root(current_handle)) {
         ancestors.emplace_back(distance_index.start_end_traversal_of(current_handle));
-        current_handle = distance_index.get_parent(current_handle);
+        net_handle_t parent_handle = distance_index.get_parent(current_handle);
+        crash_unless(parent_handle != current_handle);
+        current_handle = parent_handle;
     }
 
-
     //Now add the root-level snarl or chain
     if (distance_index.is_root_snarl(current_handle)) {
         //First thing is a snarl, so add the snarl's connected component number
@@ -1065,11 +1068,7 @@ ZipCode::snarl_code_t ZipCode::get_regular_snarl_code(const net_handle_t& snarl,
     snarl_code.set_code_type(1);
 
     //The number of children
-    size_t child_count = 0;
-    distance_index.for_each_child(snarl, [&] (const net_handle_t& child) {
-        child_count++;
-    });
-    snarl_code.set_child_count(child_count);
+    snarl_code.set_child_count(distance_index.get_snarl_child_count(snarl));
 
     //Chain prefix sum value for the start of the snarl, which is the prefix sum of the start node + length of the start node
     net_handle_t start_node = distance_index.get_node_from_sentinel(distance_index.get_bound(snarl, false, false));
@@ -1100,11 +1099,7 @@ ZipCode::snarl_code_t ZipCode::get_irregular_snarl_code(const net_handle_t& snar
     snarl_code.set_code_type(distance_index.is_dag(snarl) ? 0 : 2);
 
     //The number of children
-    size_t child_count = 0;
-    distance_index.for_each_child(snarl, [&] (const net_handle_t& child) {
-        child_count++;
-    });
-    snarl_code.set_child_count(child_count);
+    snarl_code.set_child_count(distance_index.get_snarl_child_count(snarl));
 
     //Chain prefix sum value for the start of the snarl, which is the prefix sum of the start node + length of the start node
     net_handle_t start_node = distance_index.get_node_from_sentinel(distance_index.get_bound(snarl, false, false));