From f6a295a28dcea463044da252853fd78c5eab5929 Mon Sep 17 00:00:00 2001 From: AndreaGuarracino Date: Thu, 23 Jul 2026 15:53:31 +0200 Subject: [PATCH] Speed up transitive closure and GFA emission Share the per-chunk seen-bitvector across BFS workers by reference instead of deep-copying it per worker. The copy was quadratic in input length (input x chunks x workers) and dominated large runs: on a 392 Mbp yeast pangenome (182 Mbp graph, -t8) it cuts runtime 895s -> 98s (9.1x), output byte-identical. Smaller output-preserving wins: - range_buffer uses FxHashMap instead of SipHash (per-base hot map) - skip last_seq_pos insert when min-repeat-distance is 0 - flush_ranges uses in-place retain (no temp Vec) - uppercase sequences in place on import - stream S-lines and format ints with itoa in GFA emission - byte-based CIGAR parser Byte-identical GFA across HLA (28 loci) and yeast pangenomes up to 182 Mbp. 116 tests pass. --- Cargo.lock | 8 +++++ Cargo.toml | 2 ++ src/cigar.rs | 41 +++++++++---------------- src/gfa.rs | 74 +++++++++++++++++++++++++-------------------- src/seqindex.rs | 8 ++--- src/transclosure.rs | 58 +++++++++++++++++++++++------------ 6 files changed, 109 insertions(+), 82 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a3fd3275..47bb040f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -808,6 +808,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "rustix" version = "1.1.4" @@ -853,6 +859,7 @@ dependencies = [ "flate2", "fm-index", "iitree-rs", + "itoa", "libc", "memmap2", "once_cell", @@ -860,6 +867,7 @@ dependencies = [ "portable-atomic", "rayon", "rdst", + "rustc-hash", "sucds", "uf_rush", ] diff --git a/Cargo.toml b/Cargo.toml index 8d422859..903cbfdb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,8 @@ sucds = "0.8" # Succinct data structures (rank/select) uf_rush = "0.1" # Lock-free union-find rdst = "0.20" # Fast parallel radix sort portable-atomic = "1.9" # 128-bit atomics for dset64 +rustc-hash = "2" # Fast non-cryptographic hasher (FxHashMap) for integer-keyed hot maps +itoa = "1" # Fast integer-to-decimal formatting for GFA emission # CLI dependencies clap = { version = "4.5", features = ["derive"] } # Command-line argument parsing diff --git a/src/cigar.rs b/src/cigar.rs index 6e648d8c..e913f40d 100644 --- a/src/cigar.rs +++ b/src/cigar.rs @@ -11,32 +11,21 @@ pub struct CigarOp { /// CIGAR format: alternating runs of digits and operation characters /// Example: "10M2I5M" -> [(10, 'M'), (2, 'I'), (5, 'M')] pub fn cigar_from_string(s: &str) -> Vec { - let mut cigar = Vec::new(); - let mut number = String::new(); - let mut op_type: u8 = 0; - - for c in s.chars() { - if c.is_ascii_digit() { - if op_type == 0 { - number.push(c); - } else { - // We have a complete operation - if let Ok(len) = number.parse::() { - cigar.push(CigarOp { len, op: op_type }); - } - number.clear(); - op_type = 0; - number.push(c); - } - } else { - op_type = c as u8; - } - } - - // Handle final operation - if !number.is_empty() && op_type != 0 { - if let Ok(len) = number.parse::() { - cigar.push(CigarOp { len, op: op_type }); + // Parse directly over ASCII bytes with an integer accumulator: no String + // buffer, no per-op reparse. A CIGAR op char follows its digit run, so the + // op is pushed exactly when its op byte is seen (no trailing flush needed). + let mut cigar = Vec::with_capacity(s.len() / 2); + let mut len: u64 = 0; + let mut have_digit = false; + + for &b in s.as_bytes() { + if b.is_ascii_digit() { + len = len * 10 + (b - b'0') as u64; + have_digit = true; + } else if have_digit { + cigar.push(CigarOp { len, op: b }); + len = 0; + have_digit = false; } } diff --git a/src/gfa.rs b/src/gfa.rs index 3fb1fc40..0648176f 100644 --- a/src/gfa.rs +++ b/src/gfa.rs @@ -57,41 +57,45 @@ pub fn emit_gfa( // The graph sequence has `seq_v_slice.len()` bytes; use that as the end // sentinel for the last node (the bitvector has size graph_length + 1). let graph_len = seq_v_slice.len(); - let node_sequences: Vec<(usize, String)> = (1..=n_nodes) - .map(|id| { - let node_start = match seq_id_cbv.select(id) { - Some(pos) => pos, - None => return (id, String::new()), - }; - let node_end = seq_id_cbv.select(id + 1).unwrap_or(graph_len); - let node_length = node_end - node_start; - if node_start + node_length > graph_len { - return (id, String::new()); - } - let seq = &seq_v_slice[node_start..node_start + node_length]; - let seq_string = String::from_utf8_lossy(seq).to_string(); - (id, seq_string) - }) - .collect(); - // Write node records in order - for (id, seq) in node_sequences { - if !seq.is_empty() { - writeln!(out, "S\t{id}\t{seq}")?; + // Reusable line/int buffers to avoid per-record allocation and fmt dispatch. + let mut line: Vec = Vec::new(); + let mut ib = itoa::Buffer::new(); + + // Write nodes (S lines) by streaming raw sequence bytes (seq_v is uppercased + // ASCII DNA, so this is byte-identical to from_utf8_lossy). No per-node String. + for id in 1..=n_nodes { + let node_start = match seq_id_cbv.select(id) { + Some(pos) => pos, + None => continue, + }; + let node_end = seq_id_cbv.select(id + 1).unwrap_or(graph_len); + let node_length = node_end - node_start; + if node_length == 0 || node_start + node_length > graph_len { + continue; } + let seq = &seq_v_slice[node_start..node_start + node_length]; + out.write_all(b"S\t")?; + out.write_all(ib.format(id).as_bytes())?; + out.write_all(b"\t")?; + out.write_all(seq)?; + out.write_all(b"\n")?; } // Write links (L lines) for (from, to) in links { if *from != 0 && *to != 0 { - writeln!( - out, - "L\t{}\t{}\t{}\t{}\t0M", - offset(*from), - if is_rev(*from) { "-" } else { "+" }, - offset(*to), - if is_rev(*to) { "-" } else { "+" } - )?; + line.clear(); + line.extend_from_slice(b"L\t"); + line.extend_from_slice(ib.format(offset(*from)).as_bytes()); + line.extend_from_slice(if is_rev(*from) { b"\t-\t" } else { b"\t+\t" }); + line.extend_from_slice(ib.format(offset(*to)).as_bytes()); + line.extend_from_slice(if is_rev(*to) { + b"\t-\t0M\n" + } else { + b"\t+\t0M\n" + }); + out.write_all(&line)?; } } @@ -254,15 +258,19 @@ pub fn emit_gfa( // Write path let seq_name = seqidx.nth_name(i).unwrap_or_else(|| format!("seq{i}")); - write!(out, "P\t{seq_name}\t")?; - + line.clear(); + line.extend_from_slice(b"P\t"); + line.extend_from_slice(seq_name.as_bytes()); + line.push(b'\t'); for (idx, p) in path_v.iter().enumerate() { if idx > 0 { - write!(out, ",")?; + line.push(b','); } - write!(out, "{}{}", offset(*p), if is_rev(*p) { "-" } else { "+" })?; + line.extend_from_slice(ib.format(offset(*p)).as_bytes()); + line.push(if is_rev(*p) { b'-' } else { b'+' }); } - writeln!(out, "\t*")?; + line.extend_from_slice(b"\t*\n"); + out.write_all(&line)?; } // Cleanup mmap (automatic via Drop) diff --git a/src/seqindex.rs b/src/seqindex.rs index f06f3ecd..e5147743 100644 --- a/src/seqindex.rs +++ b/src/seqindex.rs @@ -226,13 +226,13 @@ impl SeqIndex { // Record sequence boundary seq_boundary_positions.push(seq_bytes_written); - // Write upper-case sequence - let seq_upper = seq.to_uppercase(); + // Write upper-case sequence (ASCII DNA: in-place, no alloc) + seq.make_ascii_uppercase(); seq_out - .write_all(seq_upper.as_bytes()) + .write_all(seq.as_bytes()) .map_err(|e| format!("Failed to write sequence: {e}"))?; - seq_bytes_written += seq_upper.len() as u64; + seq_bytes_written += seq.len() as u64; self.seq_count += 1; // Check EOF diff --git a/src/transclosure.rs b/src/transclosure.rs index e4228f26..48caaf58 100644 --- a/src/transclosure.rs +++ b/src/transclosure.rs @@ -4,6 +4,8 @@ // identifying equivalence classes that form nodes in the variation graph. use std::collections::HashMap; + +use rustc_hash::FxHashMap; use std::io; use std::sync::atomic::Ordering; use std::sync::{Arc, RwLock}; @@ -113,7 +115,7 @@ impl AtomicBitVec { pub fn extend_range( s_pos: u64, q_pos: PosT, - range_buffer: &mut HashMap, + range_buffer: &mut FxHashMap, seqidx: &SeqIndex, node_iitree: &mut AdaptiveTree, path_iitree: &mut AdaptiveTree, @@ -205,19 +207,28 @@ fn flush_single_range( /// Flush all ranges in the buffer that aren't at s_pos pub fn flush_ranges( s_pos: u64, - range_buffer: &mut HashMap, + range_buffer: &mut FxHashMap, node_iitree: &mut AdaptiveTree, path_iitree: &mut AdaptiveTree, ) -> io::Result<()> { - let to_flush: Vec<_> = range_buffer - .iter() - .filter(|(_, range)| range.end != s_pos) - .map(|(k, v)| (*k, *v)) - .collect(); - - for (key, range) in to_flush { - flush_single_range(&range, key, node_iitree, path_iitree)?; - range_buffer.remove(&key); + // Flush in place: same iteration order and same set of removals as + // collect-then-remove, but without the temporary Vec or the extra + // per-key hash lookups. Capture the first error and stop flushing. + let mut err: Option = None; + range_buffer.retain(|&key, range| { + if range.end != s_pos { + if err.is_none() { + if let Err(e) = flush_single_range(range, key, node_iitree, path_iitree) { + err = Some(e); + } + } + false + } else { + true + } + }); + if let Some(e) = err { + return Err(e); } Ok(()) @@ -325,7 +336,7 @@ fn write_graph_chunk( node_iitree: &mut AdaptiveTree, path_iitree: &mut AdaptiveTree, seq_v_out: &mut Vec, - range_buffer: &mut HashMap, + range_buffer: &mut FxHashMap, dsets: Vec<(u64, u64)>, repeat_max: u64, min_repeat_dist: u64, @@ -414,7 +425,11 @@ fn write_graph_chunk( } else { todos.entry(curr_seq_count).or_default().push(curr_q_pos); } - last_seq_pos.insert(curr_seq_id, curr_q_pos); + // last_seq_pos is only read by close_to_prev, which only runs when + // min_repeat_dist != 0. Skip the per-base insert otherwise. + if min_repeat_dist != 0 { + last_seq_pos.insert(curr_seq_id, curr_q_pos); + } } } @@ -586,14 +601,15 @@ pub fn compute_transitive_closures( let mut q_seen_bv = vec![false; input_seq_length]; // Range buffer for writing to iitrees - let mut range_buffer = Some(HashMap::new()); + let mut range_buffer = Some(FxHashMap::default()); let mut bases_seen = 0u64; // Writer thread handle for pipelining (like C++) // The thread writes while we compute the next batch - let mut writer_thread: Option, HashMap)>>> = - None; + let mut writer_thread: Option< + thread::JoinHandle, FxHashMap)>>, + > = None; // Main loop: process input sequence in chunks let mut i = 0; @@ -695,13 +711,17 @@ pub fn compute_transitive_closures( } }); + // Workers only read q_seen_bv during the scope (it is mutated after the + // scope joins), so share one immutable slice instead of deep-copying the + // whole bitvector per worker. At scale the per-worker clone was the + // dominant cost (chunks x workers x input_length bytes of memcpy). + let q_seen_bv_ref: &[bool] = &q_seen_bv; rayon::scope(|s| { for _worker_idx in 0..(num_threads * 2) { let todo_out = Arc::clone(&todo_out); let todo_in = Arc::clone(&todo_in); let aln_iitree = Arc::clone(&aln_iitree_clone); let q_curr_bv = Arc::clone(&q_curr_bv_shared); - let q_seen_bv_clone = q_seen_bv.clone(); let seqidx = Arc::clone(&seqidx_clone); let pending_clone = Arc::clone(&pending); @@ -716,7 +736,7 @@ pub fn compute_transitive_closures( explore_overlaps_discovery( &Match::new(n, n + match_len, pos), - &q_seen_bv_clone, + q_seen_bv_ref, &q_curr_bv, &aln_iitree, &todo_in, @@ -1096,7 +1116,7 @@ pub fn compute_transitive_closures( let seqidx_clone = Arc::clone(&seqidx); writer_thread = Some(thread::spawn( - move || -> io::Result<(Vec, HashMap)> { + move || -> io::Result<(Vec, FxHashMap)> { let mut node_guard = node_iitree_clone.write().unwrap(); let mut path_guard = path_iitree_clone.write().unwrap();