Auto merge of #139758 - Zoxc:thread-local-graph, r=oli-obk

Use thread local dep graph encoding

This adds thread local encoding of dep graph nodes. Each thread has a `MemEncoder` that gets flushed to the global `FileEncoder` when it exceeds 64 kB. Each thread also has a local cache of dep indices. This means there can now be empty gaps in `SerializedDepGraph`.

Indices are marked green and also allocated by the new atomic operation `DepNodeColorMap::try_mark_green` as the encoder lock is removed.

Author: bors

compiler/rustc_data_structures/src/sync.rs

@@ -43,7 +43,7 @@ pub use self::freeze::{FreezeLock, FreezeReadGuard, FreezeWriteGuard};
 pub use self::lock::{Lock, LockGuard, Mode};
 pub use self::mode::{is_dyn_thread_safe, set_dyn_thread_safe_mode};
 pub use self::parallel::{
-    join, par_for_each_in, par_map, parallel_guard, scope, spawn, try_par_for_each_in,
+    broadcast, join, par_for_each_in, par_map, parallel_guard, scope, spawn, try_par_for_each_in,
 };
 pub use self::vec::{AppendOnlyIndexVec, AppendOnlyVec};
 pub use self::worker_local::{Registry, WorkerLocal};

compiler/rustc_data_structures/src/sync/parallel.rs

@@ -237,3 +237,13 @@ pub fn par_map<I: DynSend, T: IntoIterator<Item = I>, R: DynSend, C: FromIterato
         }
     })
 }
+
+pub fn broadcast<R: DynSend>(op: impl Fn(usize) -> R + DynSync) -> Vec<R> {
+    if mode::is_dyn_thread_safe() {
+        let op = FromDyn::from(op);
+        let results = rayon_core::broadcast(|context| op.derive(op(context.index())));
+        results.into_iter().map(|r| r.into_inner()).collect()
+    } else {
+        vec![op(0)]
+    }
+}

compiler/rustc_incremental/src/persist/save.rs

@@ -44,10 +44,6 @@ pub(crate) fn save_dep_graph(tcx: TyCtxt<'_>) {
         sess.time("assert_dep_graph", || assert_dep_graph(tcx));
         sess.time("check_dirty_clean", || dirty_clean::check_dirty_clean_annotations(tcx));
 
-        if sess.opts.unstable_opts.incremental_info {
-            tcx.dep_graph.print_incremental_info()
-        }
-
         join(
             move || {
                 sess.time("incr_comp_persist_dep_graph", || {
@@ -172,12 +168,5 @@ pub(crate) fn build_dep_graph(
     // First encode the commandline arguments hash
     sess.opts.dep_tracking_hash(false).encode(&mut encoder);
 
-    Some(DepGraph::new(
-        sess,
-        prev_graph,
-        prev_work_products,
-        encoder,
-        sess.opts.unstable_opts.query_dep_graph,
-        sess.opts.unstable_opts.incremental_info,
-    ))
+    Some(DepGraph::new(sess, prev_graph, prev_work_products, encoder))
 }

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -11,7 +11,7 @@ use rustc_data_structures::outline;
 use rustc_data_structures::profiling::QueryInvocationId;
 use rustc_data_structures::sharded::{self, ShardedHashMap};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
-use rustc_data_structures::sync::{AtomicU64, Lock};
+use rustc_data_structures::sync::{AtomicU64, Lock, is_dyn_thread_safe};
 use rustc_data_structures::unord::UnordMap;
 use rustc_errors::DiagInner;
 use rustc_index::IndexVec;
@@ -124,19 +124,11 @@ impl<D: Deps> DepGraph<D> {
         prev_graph: Arc<SerializedDepGraph>,
         prev_work_products: WorkProductMap,
         encoder: FileEncoder,
-        record_graph: bool,
-        record_stats: bool,
     ) -> DepGraph<D> {
         let prev_graph_node_count = prev_graph.node_count();
 
-        let current = CurrentDepGraph::new(
-            session,
-            prev_graph_node_count,
-            encoder,
-            record_graph,
-            record_stats,
-            Arc::clone(&prev_graph),
-        );
+        let current =
+            CurrentDepGraph::new(session, prev_graph_node_count, encoder, Arc::clone(&prev_graph));
 
         let colors = DepNodeColorMap::new(prev_graph_node_count);
 
@@ -1052,17 +1044,8 @@ impl<D: Deps> DepGraph<D> {
         }
     }
 
-    pub fn print_incremental_info(&self) {
-        if let Some(data) = &self.data {
-            data.current.encoder.print_incremental_info(
-                data.current.total_read_count.load(Ordering::Relaxed),
-                data.current.total_duplicate_read_count.load(Ordering::Relaxed),
-            )
-        }
-    }
-
     pub fn finish_encoding(&self) -> FileEncodeResult {
-        if let Some(data) = &self.data { data.current.encoder.finish() } else { Ok(0) }
+        if let Some(data) = &self.data { data.current.encoder.finish(&data.current) } else { Ok(0) }
     }
 
     pub(crate) fn next_virtual_depnode_index(&self) -> DepNodeIndex {
@@ -1179,17 +1162,15 @@ pub(super) struct CurrentDepGraph<D: Deps> {
 
     /// These are simple counters that are for profiling and
     /// debugging and only active with `debug_assertions`.
-    total_read_count: AtomicU64,
-    total_duplicate_read_count: AtomicU64,
+    pub(super) total_read_count: AtomicU64,
+    pub(super) total_duplicate_read_count: AtomicU64,
 }
 
 impl<D: Deps> CurrentDepGraph<D> {
     fn new(
         session: &Session,
         prev_graph_node_count: usize,
         encoder: FileEncoder,
-        record_graph: bool,
-        record_stats: bool,
         previous: Arc<SerializedDepGraph>,
     ) -> Self {
         let mut stable_hasher = StableHasher::new();
@@ -1211,14 +1192,7 @@ impl<D: Deps> CurrentDepGraph<D> {
             session.opts.unstable_opts.incremental_verify_ich || cfg!(debug_assertions);
 
         CurrentDepGraph {
-            encoder: GraphEncoder::new(
-                encoder,
-                prev_graph_node_count,
-                record_graph,
-                record_stats,
-                &session.prof,
-                previous,
-            ),
+            encoder: GraphEncoder::new(session, encoder, prev_graph_node_count, previous),
             anon_node_to_index: ShardedHashMap::with_capacity(
                 // FIXME: The count estimate is off as anon nodes are only a portion of the nodes.
                 new_node_count_estimate / sharded::shards(),
@@ -1345,6 +1319,7 @@ impl Default for TaskDeps {
 // array, using one u32 per entry.
 pub(super) struct DepNodeColorMap {
     values: IndexVec<SerializedDepNodeIndex, AtomicU32>,
+    sync: bool,
 }
 
 const COMPRESSED_NONE: u32 = u32::MAX;
@@ -1353,7 +1328,10 @@ const COMPRESSED_RED: u32 = u32::MAX - 1;
 impl DepNodeColorMap {
     fn new(size: usize) -> DepNodeColorMap {
         debug_assert!(COMPRESSED_RED > DepNodeIndex::MAX_AS_U32);
-        DepNodeColorMap { values: (0..size).map(|_| AtomicU32::new(COMPRESSED_NONE)).collect() }
+        DepNodeColorMap {
+            values: (0..size).map(|_| AtomicU32::new(COMPRESSED_NONE)).collect(),
+            sync: is_dyn_thread_safe(),
+        }
     }
 
     #[inline]
@@ -1362,6 +1340,37 @@ impl DepNodeColorMap {
         if value <= DepNodeIndex::MAX_AS_U32 { Some(DepNodeIndex::from_u32(value)) } else { None }
     }
 
+    /// This tries to atomically mark a node green and assign `index` as the new
+    /// index. This returns `Ok` if `index` gets assigned, otherwise it returns
+    /// the alreadly allocated index in `Err`.
+    #[inline]
+    pub(super) fn try_mark_green(
+        &self,
+        prev_index: SerializedDepNodeIndex,
+        index: DepNodeIndex,
+    ) -> Result<(), DepNodeIndex> {
+        let value = &self.values[prev_index];
+        if self.sync {
+            match value.compare_exchange(
+                COMPRESSED_NONE,
+                index.as_u32(),
+                Ordering::Relaxed,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => Ok(()),
+                Err(v) => Err(DepNodeIndex::from_u32(v)),
+            }
+        } else {
+            let v = value.load(Ordering::Relaxed);
+            if v == COMPRESSED_NONE {
+                value.store(index.as_u32(), Ordering::Relaxed);
+                Ok(())
+            } else {
+                Err(DepNodeIndex::from_u32(v))
+            }
+        }
+    }
+
     #[inline]
     pub(super) fn get(&self, index: SerializedDepNodeIndex) -> Option<DepNodeColor> {
         match self.values[index].load(Ordering::Acquire) {

compiler/rustc_query_system/src/dep_graph/mod.rs

@@ -12,6 +12,7 @@ pub(crate) use graph::DepGraphData;
 pub use graph::{DepGraph, DepNodeIndex, TaskDepsRef, WorkProduct, WorkProductMap, hash_result};
 pub use query::DepGraphQuery;
 use rustc_data_structures::profiling::SelfProfilerRef;
+use rustc_data_structures::sync::DynSync;
 use rustc_session::Session;
 pub use serialized::{SerializedDepGraph, SerializedDepNodeIndex};
 use tracing::instrument;
@@ -89,7 +90,7 @@ pub trait DepContext: Copy {
     }
 }
 
-pub trait Deps {
+pub trait Deps: DynSync {
     /// Execute the operation with provided dependencies.
     fn with_deps<OP, R>(deps: TaskDepsRef<'_>, op: OP) -> R
     where