diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index f537e6f91..fa4deee73 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -5,16 +5,20 @@ use crate::edge_unpacker::EdgeUnpacker; use crate::errors::{EngineError, Result}; use crate::output::OutputHandler; -use crate::state::{LookupKey, NodeRecord, RawEdge, ResolvedProjection, SamplingState}; +use crate::state::{ + CommitPart, LookupKey, NodeRecord, RawEdge, ResolvedProjection, SamplingState, TypeState, +}; use diskgraph_core::sampling::*; use diskgraph_core::sortkey::SortKeyDecoder; use diskgraph_core::tt_builder::{Value, WEIGHT_COL}; use diskgraph_core::time::{abs_from_unit, Seconds, TimeUnit}; use diskgraph_index::layout::{EdgeListView, RelationSegment}; use diskgraph_index::reader::ShardReader; +use gxhash::HashMap as GxHashMap; use rayon::prelude::*; use std::collections::HashMap; use std::path::Path; +use std::sync::atomic::{AtomicU32, Ordering::Relaxed}; use std::sync::Arc; use std::time::Instant; use tracing::{debug, info, warn}; @@ -464,11 +468,38 @@ pub struct DiskGraphEngine { pub thread_pool: rayon::ThreadPool, /// Record edges through the global `index_map`; same output. pub legacy_recording: bool, + /// Smallest layer (in edges) whose per-seed dedup runs on the pool. + pub parallel_commit_min_edges: usize, } const NEW_NODE: u32 = u32::MAX; const PENDING: u32 = 1 << 31; +fn dedup_seeds( + maps: &mut [GxHashMap], + order: &[(u64, u32)], + bounds: &[u32], + res: &[AtomicU32], + part: &mut CommitPart, +) { + part.new_edges.clear(); + part.new_bounds.clear(); + for (s, map) in maps.iter_mut().enumerate() { + for &(key, i) in &order[bounds[s] as usize..bounds[s + 1] as usize] { + let r = match map.entry(key) { + std::collections::hash_map::Entry::Occupied(o) => *o.get(), + std::collections::hash_map::Entry::Vacant(v) => { + v.insert(PENDING | i); + part.new_edges.push((key, i)); + NEW_NODE + } + }; + res[i as usize].store(r, Relaxed); + } + part.new_bounds.push(part.new_edges.len() as u32); + } +} + #[inline] fn node_key(type_id: usize, global_id: u64) -> u64 { debug_assert!(global_id < 1 << 56); @@ -860,6 +891,7 @@ impl DiskGraphEngine { random_seed, thread_pool, legacy_recording: false, + parallel_commit_min_edges: 1 << 15, } } @@ -2189,6 +2221,17 @@ impl DiskGraphEngine { // 3. Process Layers (BFS) let compiled_plan = self.compile_plan(&request.plan)?; + let anchors: Vec> = self + .id_to_node_type + .iter() + .map(|t| { + request + .seed_timestamps + .as_ref() + .and_then(|m| m.get(t)) + .map(Vec::as_slice) + }) + .collect(); for layer_idx in 0..request.plan.layers.len() { let hop_start = Instant::now(); let edges_before = state.stats.total_edges; @@ -2198,6 +2241,7 @@ impl DiskGraphEngine { layer_idx, &compiled_plan.layers[layer_idx], request, + &anchors, state, output, cutoff, @@ -2549,7 +2593,7 @@ impl DiskGraphEngine { rel_id: usize, segment: &RelationSegment<'_>, fanout_spec: &diskgraph_core::sampling::FanoutSpec, - request: &BatchSamplingRequest, + anchors: &[Option<&[i64]>], dst_scratch: &mut Vec, key_scratch: &mut Vec, cutoff: Option<&NodeCutoff<'_>>, @@ -2584,16 +2628,11 @@ impl DiskGraphEngine { // root seed type — not an arbitrary `keys().next()` type, which // is wrong (and out-of-bounds) once seed arrays differ in length // across types (e.g. LP lhs vs rhs + negatives). - let anchor = self - .id_to_node_type + let anchor = anchors .get(node.seed_type_id as usize) - .and_then(|seed_type| { - request - .seed_timestamps - .as_ref() - .and_then(|m| m.get(seed_type)) - .and_then(|v| v.get(node.seed_index as usize).copied()) - }); + .copied() + .flatten() + .and_then(|v| v.get(node.seed_index as usize).copied()); if let (Some(dec), Some(anch)) = (decoder, anchor) { let newer = unpacker .partition_point(|key| dec.decode(key).0.is_some_and(|ts| ts > anch)); @@ -2829,6 +2868,7 @@ impl DiskGraphEngine { layer_idx: usize, layer_spec: &CompiledLayerSpec, request: &BatchSamplingRequest, + anchors: &[Option<&[i64]>], state: &mut SamplingState, output: &mut impl OutputHandler, cutoff: Option<&NodeCutoff<'_>>, @@ -2880,7 +2920,15 @@ impl DiskGraphEngine { let segments = edges_view.relation_segments().unwrap_or_default(); self.process_node_relations( - layer_idx, src_type_id, &segments, node, layer_spec, request, state, output, + layer_idx, + src_type_id, + &segments, + node, + layer_spec, + request, + anchors, + state, + output, cutoff, )?; } @@ -2974,7 +3022,7 @@ impl DiskGraphEngine { rel_id, segment, fanout_spec, - request, + anchors, &mut dst_scratch, &mut key_scratch, cutoff, @@ -3040,6 +3088,7 @@ impl DiskGraphEngine { node: &NodeRecord, layer_spec: &CompiledLayerSpec, request: &BatchSamplingRequest, + anchors: &[Option<&[i64]>], state: &mut SamplingState, output: &mut impl OutputHandler, cutoff: Option<&NodeCutoff<'_>>, @@ -3076,7 +3125,7 @@ impl DiskGraphEngine { rel_id, segment, fanout_spec, - request, + anchors, &mut dst_scratch, &mut key_scratch, cutoff, @@ -3131,28 +3180,47 @@ impl DiskGraphEngine { *cur += 1; } - cs.res.resize(n, 0); - cs.new_edges.clear(); - cs.new_bounds.clear(); - for s in 0..slots { - let map = &mut state.seed_maps[s]; - for &(key, i) in &cs.order[cs.bounds[s] as usize..cs.bounds[s + 1] as usize] { - match map.entry(key) { - std::collections::hash_map::Entry::Occupied(o) => cs.res[i as usize] = *o.get(), - std::collections::hash_map::Entry::Vacant(v) => { - v.insert(PENDING | i); - cs.res[i as usize] = NEW_NODE; - cs.new_edges.push((key, i)); - } - } - } - cs.new_bounds.push(cs.new_edges.len() as u32); + if cs.res.len() < n { + cs.res.resize_with(n, Default::default); } + let threads = self.thread_pool.current_num_threads(); + // 32k default: at 8 threads, layers of a few thousand to ~25k edges ran + // 10-50% slower on the pool, layers of 100k+ edges faster. Cuts fall on + // seed boundaries, so parts are uneven; 4 per thread lets the pool balance. + let parts = if threads > 1 && n >= self.parallel_commit_min_edges { + threads * 4 + } else { + 1 + }; + let mut cuts = vec![0]; + for k in 1..parts { + let cut = cs + .bounds + .partition_point(|&b| (b as usize) < n * k / parts) + .min(slots); + if cut > *cuts.last().unwrap() { + cuts.push(cut); + } + } + if *cuts.last().unwrap() < slots { + cuts.push(slots); + } + cs.parts.resize_with(cuts.len() - 1, Default::default); + self.for_each_part( + &mut state.seed_maps, + &cuts, + &mut cs.parts, + |maps, part, range| dedup_seeds(maps, &cs.order, &cs.bounds[range], &cs.res, part), + ); + let mut dst_states: Vec> = Vec::new(); + dst_states.resize_with(self.readers.len(), || None); for (i, e) in chunks.iter().flatten().enumerate() { - let r = cs.res[i]; + let r = cs.res[i].load(Relaxed); let dst_local_id = if r == NEW_NODE { - let ts = state.type_states.entry(e.dst_type_id).or_default(); + let ts = dst_states[e.dst_type_id].get_or_insert_with(|| { + state.type_states.remove(&e.dst_type_id).unwrap_or_default() + }); if ts.nodes.len() >= PENDING as usize { return Err(EngineError::Internal(format!( "too many nodes of type {}", @@ -3169,10 +3237,10 @@ impl DiskGraphEngine { ts.nodes.push(node); ts.next_frontier.push(node); state.stats.record_node(layer_idx + 1, e.dst_type_id); - cs.res[i] = id; + cs.res[i].store(id, Relaxed); id } else if r & PENDING != 0 { - cs.res[(r & !PENDING) as usize] + cs.res[(r & !PENDING) as usize].load(Relaxed) } else { r }; @@ -3186,19 +3254,55 @@ impl DiskGraphEngine { output, )?; } - - let mut lo = 0; - for (s, &hi) in cs.new_bounds.iter().enumerate() { - let map = &mut state.seed_maps[s]; - for &(key, i) in &cs.new_edges[lo..hi as usize] { - map.insert(key, cs.res[i as usize]); + for (t, ts) in dst_states.into_iter().enumerate() { + if let Some(ts) = ts { + state.type_states.insert(t, ts); } - lo = hi as usize; } + + self.for_each_part( + &mut state.seed_maps, + &cuts, + &mut cs.parts, + |maps, part, _| { + let mut lo = 0; + for (map, &hi) in maps.iter_mut().zip(&part.new_bounds) { + for &(key, i) in &part.new_edges[lo..hi as usize] { + map.insert(key, cs.res[i as usize].load(Relaxed)); + } + lo = hi as usize; + } + }, + ); state.commit_scratch = cs; Ok(()) } + fn for_each_part( + &self, + maps: &mut [GxHashMap], + cuts: &[usize], + parts: &mut [CommitPart], + f: F, + ) where + F: Fn(&mut [GxHashMap], &mut CommitPart, std::ops::RangeInclusive) + Sync, + { + let mut rest = maps; + let mut jobs = Vec::with_capacity(parts.len()); + for (w, part) in cuts.windows(2).zip(parts.iter_mut()) { + let (head, tail) = std::mem::take(&mut rest).split_at_mut(w[1] - w[0]); + rest = tail; + jobs.push((head, part, w[0]..=w[1])); + } + if jobs.len() == 1 { + let (m, p, r) = jobs.pop().unwrap(); + f(m, p, r); + } else { + self.thread_pool + .install(|| jobs.into_par_iter().for_each(|(m, p, r)| f(m, p, r))); + } + } + #[allow(clippy::too_many_arguments)] fn emit_edge( &self, diff --git a/packages/diskgraph/diskgraph-engine/src/pseudo_time_tests.rs b/packages/diskgraph/diskgraph-engine/src/pseudo_time_tests.rs index 429e76ecf..0f77a7f1f 100644 --- a/packages/diskgraph/diskgraph-engine/src/pseudo_time_tests.rs +++ b/packages/diskgraph/diskgraph-engine/src/pseudo_time_tests.rs @@ -3,7 +3,6 @@ // SPDX-License-Identifier: Apache-2.0 use super::*; -use diskgraph_core::dataset::DatasetId; use diskgraph_index::layout::RelationLayout; fn sample(times: &[u8], fanout: u32, seed: u64, mode: SamplingMode) -> Result> { @@ -37,22 +36,6 @@ fn sample(times: &[u8], fanout: u32, seed: u64, mode: SamplingMode) -> Result Result, pub cursor: Vec, pub order: Vec<(u64, u32)>, - pub res: Vec, + pub res: Vec, + pub parts: Vec, +} + +#[derive(Default)] +pub struct CommitPart { pub new_edges: Vec<(u64, u32)>, pub new_bounds: Vec, } @@ -337,8 +342,11 @@ impl SamplingState { clear_and_cap(&mut cs.cursor, RESET_THRESHOLD); clear_and_cap(&mut cs.order, RESET_THRESHOLD); clear_and_cap(&mut cs.res, RESET_THRESHOLD); - clear_and_cap(&mut cs.new_edges, RESET_THRESHOLD); - clear_and_cap(&mut cs.new_bounds, RESET_THRESHOLD); + let per_part = RESET_THRESHOLD / cs.parts.len().max(1); + for part in &mut cs.parts { + clear_and_cap(&mut part.new_edges, per_part); + clear_and_cap(&mut part.new_bounds, per_part); + } self.seed_maps.truncate(SEED_MAPS_KEEP); let mut budget = RESET_THRESHOLD; for m in &mut self.seed_maps { @@ -434,8 +442,13 @@ mod tests { .collect(); let cs = &mut state.commit_scratch; cs.order.reserve(RESET_THRESHOLD * 4); - cs.res.reserve(RESET_THRESHOLD * 4); - cs.new_edges.reserve(RESET_THRESHOLD * 4); + cs.res.resize_with(RESET_THRESHOLD * 4, Default::default); + cs.parts = (0..4) + .map(|_| CommitPart { + new_edges: Vec::with_capacity(RESET_THRESHOLD * 2), + new_bounds: Vec::new(), + }) + .collect(); state.seed_maps = (0..64) .map(|_| GxHashMap::with_capacity(RESET_THRESHOLD / 16)) .collect(); @@ -446,7 +459,13 @@ mod tests { let cs = &state.commit_scratch; assert!(cs.order.capacity() <= RESET_THRESHOLD); assert!(cs.res.capacity() <= RESET_THRESHOLD); - assert!(cs.new_edges.capacity() <= RESET_THRESHOLD); + assert!( + cs.parts + .iter() + .map(|p| p.new_edges.capacity()) + .sum::() + <= RESET_THRESHOLD + ); let map_capacity: usize = state.seed_maps.iter().map(GxHashMap::capacity).sum(); assert!(map_capacity <= RESET_THRESHOLD); assert!(map_capacity > 0); diff --git a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs index 9f416b487..12f474989 100644 --- a/packages/diskgraph/diskgraph-engine/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-engine/src/tests_unit.rs @@ -2618,9 +2618,12 @@ fn assert_per_seed_matches_legacy( }; sampler.legacy_recording = false; let per_seed = run(sampler); + sampler.parallel_commit_min_edges = 1; + let parallel = run(sampler); sampler.legacy_recording = true; let legacy = run(sampler); assert_eq!(per_seed, legacy); + assert_eq!(parallel, legacy); per_seed }