From cfbe4bf519b297fac42b8a5f60279a812487f172 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Wed, 16 Aug 2023 17:50:40 +0200 Subject: [PATCH 1/2] Optimize hash map operations in the query system --- Cargo.lock | 3 + compiler/rustc_data_structures/Cargo.toml | 5 ++ compiler/rustc_data_structures/src/marker.rs | 1 + compiler/rustc_middle/src/query/plumbing.rs | 3 +- compiler/rustc_query_system/Cargo.toml | 6 ++ .../rustc_query_system/src/query/caches.rs | 32 +++++----- .../rustc_query_system/src/query/plumbing.rs | 63 ++++++++++++------- 7 files changed, 73 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4cb2a584e6ab1..19fa4ff1dd44a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3536,6 +3536,7 @@ dependencies = [ "cfg-if", "elsa", "ena", + "hashbrown 0.14.0", "indexmap 2.0.0", "itertools", "jobserver", @@ -4264,7 +4265,9 @@ dependencies = [ name = "rustc_query_system" version = "0.0.0" dependencies = [ + "hashbrown 0.14.0", "parking_lot 0.12.1", + "rustc-hash", "rustc-rayon-core", "rustc_ast", "rustc_data_structures", diff --git a/compiler/rustc_data_structures/Cargo.toml b/compiler/rustc_data_structures/Cargo.toml index f77bd53e76c6a..5b4d7b0c78f1b 100644 --- a/compiler/rustc_data_structures/Cargo.toml +++ b/compiler/rustc_data_structures/Cargo.toml @@ -14,6 +14,11 @@ indexmap = { version = "2.0.0" } jobserver_crate = { version = "0.1.13", package = "jobserver" } libc = "0.2" measureme = "10.0.0" +hashbrown = { version = "0.14", default-features = false, features = [ + "raw", + "inline-more", + "nightly", +] } rustc-rayon-core = { version = "0.5.0", optional = true } rustc-rayon = { version = "0.5.0", optional = true } rustc_arena = { path = "../rustc_arena" } diff --git a/compiler/rustc_data_structures/src/marker.rs b/compiler/rustc_data_structures/src/marker.rs index b067f9d4502df..a786e2265ac4c 100644 --- a/compiler/rustc_data_structures/src/marker.rs +++ b/compiler/rustc_data_structures/src/marker.rs @@ -87,6 +87,7 @@ cfg_if!( [std::sync::mpsc::Sender where T: DynSend] [std::sync::Arc where T: ?Sized + DynSync + DynSend] [std::sync::LazyLock where T: DynSend, F: DynSend] + [hashbrown::HashMap where K: DynSend, V: DynSend, S: DynSend] [std::collections::HashSet where K: DynSend, S: DynSend] [std::collections::HashMap where K: DynSend, V: DynSend, S: DynSend] [std::collections::BTreeMap where K: DynSend, V: DynSend, A: std::alloc::Allocator + Clone + DynSend] diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs index a342b5231e9a5..ea734cfc6cf36 100644 --- a/compiler/rustc_middle/src/query/plumbing.rs +++ b/compiler/rustc_middle/src/query/plumbing.rs @@ -524,7 +524,8 @@ macro_rules! define_feedable { &value, hash_result!([$($modifiers)*]), ); - cache.complete(key, erased, dep_node_index); + let key_hash = rustc_data_structures::sharded::make_hash(&key); + cache.complete(key, key_hash, erased, dep_node_index); } } } diff --git a/compiler/rustc_query_system/Cargo.toml b/compiler/rustc_query_system/Cargo.toml index 584355df80249..4856ac37c39fb 100644 --- a/compiler/rustc_query_system/Cargo.toml +++ b/compiler/rustc_query_system/Cargo.toml @@ -6,7 +6,13 @@ edition = "2021" [lib] [dependencies] +hashbrown = { version = "0.14", default-features = false, features = [ + "raw", + "inline-more", + "nightly", +] } parking_lot = "0.12" +rustc-hash = "1.1.0" rustc_ast = { path = "../rustc_ast" } rustc_data_structures = { path = "../rustc_data_structures" } rustc_errors = { path = "../rustc_errors" } diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs index 0240f012da053..d5afae700c831 100644 --- a/compiler/rustc_query_system/src/query/caches.rs +++ b/compiler/rustc_query_system/src/query/caches.rs @@ -1,10 +1,10 @@ use crate::dep_graph::DepNodeIndex; - -use rustc_data_structures::fx::FxHashMap; use rustc_data_structures::sharded::{self, Sharded}; use rustc_data_structures::sync::OnceLock; +use rustc_hash::FxHasher; use rustc_index::{Idx, IndexVec}; use std::fmt::Debug; +use std::hash::BuildHasherDefault; use std::hash::Hash; use std::marker::PhantomData; @@ -19,9 +19,9 @@ pub trait QueryCache: Sized { type Value: Copy; /// Checks if the query is already computed and in the cache. - fn lookup(&self, key: &Self::Key) -> Option<(Self::Value, DepNodeIndex)>; + fn lookup(&self, key: &Self::Key, key_hash: u64) -> Option<(Self::Value, DepNodeIndex)>; - fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex); + fn complete(&self, key: Self::Key, key_hash: u64, value: Self::Value, index: DepNodeIndex); fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)); } @@ -35,7 +35,7 @@ impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelecto } pub struct DefaultCache { - cache: Sharded>, + cache: Sharded>>, } impl Default for DefaultCache { @@ -53,8 +53,7 @@ where type Value = V; #[inline(always)] - fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> { - let key_hash = sharded::make_hash(key); + fn lookup(&self, key: &K, key_hash: u64) -> Option<(V, DepNodeIndex)> { let lock = self.cache.lock_shard_by_hash(key_hash); let result = lock.raw_entry().from_key_hashed_nocheck(key_hash, key); @@ -62,11 +61,12 @@ where } #[inline] - fn complete(&self, key: K, value: V, index: DepNodeIndex) { - let mut lock = self.cache.lock_shard_by_value(&key); - // We may be overwriting another value. This is all right, since the dep-graph - // will check that the fingerprint matches. - lock.insert(key, (value, index)); + fn complete(&self, key: K, key_hash: u64, value: V, index: DepNodeIndex) { + let mut lock = self.cache.lock_shard_by_hash(key_hash); + // Use the raw table API since we know that the + // key does not exist and we already have the hash. + lock.raw_table_mut() + .insert(key_hash, (key, (value, index)), |(k, _)| sharded::make_hash(k)); } fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) { @@ -104,12 +104,12 @@ where type Value = V; #[inline(always)] - fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> { + fn lookup(&self, _key: &(), _key_hash: u64) -> Option<(V, DepNodeIndex)> { self.cache.get().copied() } #[inline] - fn complete(&self, _key: (), value: V, index: DepNodeIndex) { + fn complete(&self, _key: (), _key_hash: u64, value: V, index: DepNodeIndex) { self.cache.set((value, index)).ok(); } @@ -147,13 +147,13 @@ where type Value = V; #[inline(always)] - fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> { + fn lookup(&self, key: &K, _key_hash: u64) -> Option<(V, DepNodeIndex)> { let lock = self.cache.lock_shard_by_hash(key.index() as u64); if let Some(Some(value)) = lock.get(*key) { Some(*value) } else { None } } #[inline] - fn complete(&self, key: K, value: V, index: DepNodeIndex) { + fn complete(&self, key: K, _key_hash: u64, value: V, index: DepNodeIndex) { let mut lock = self.cache.lock_shard_by_hash(key.index() as u64); lock.insert(key, (value, index)); } diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs index 07db15e6d8be0..4268e57cf5644 100644 --- a/compiler/rustc_query_system/src/query/plumbing.rs +++ b/compiler/rustc_query_system/src/query/plumbing.rs @@ -12,18 +12,19 @@ use crate::query::job::{report_cycle, QueryInfo, QueryJob, QueryJobId, QueryJobI use crate::query::SerializedDepNodeIndex; use crate::query::{QueryContext, QueryMap, QuerySideEffects, QueryStackFrame}; use crate::HandleCycleError; +use hashbrown::hash_map::RawEntryMut; use rustc_data_structures::fingerprint::Fingerprint; -use rustc_data_structures::fx::FxHashMap; -use rustc_data_structures::sharded::Sharded; +use rustc_data_structures::sharded::{self, Sharded}; use rustc_data_structures::stack::ensure_sufficient_stack; use rustc_data_structures::sync::Lock; #[cfg(parallel_compiler)] use rustc_data_structures::{cold_path, sync}; use rustc_errors::{DiagnosticBuilder, ErrorGuaranteed, FatalError}; +use rustc_hash::FxHasher; use rustc_span::{Span, DUMMY_SP}; use std::cell::Cell; -use std::collections::hash_map::Entry; use std::fmt::Debug; +use std::hash::BuildHasherDefault; use std::hash::Hash; use std::mem; use thin_vec::ThinVec; @@ -31,7 +32,7 @@ use thin_vec::ThinVec; use super::QueryConfig; pub struct QueryState { - active: Sharded>>, + active: Sharded, BuildHasherDefault>>, } /// Indicates the state of a query for a given key in a query map. @@ -143,7 +144,7 @@ where { /// Completes the query by updating the query cache with the `result`, /// signals the waiter and forgets the JobOwner, so it won't poison the query - fn complete(self, cache: &C, result: C::Value, dep_node_index: DepNodeIndex) + fn complete(self, cache: &C, key_hash: u64, result: C::Value, dep_node_index: DepNodeIndex) where C: QueryCache, { @@ -155,13 +156,17 @@ where // Mark as complete before we remove the job from the active state // so no other thread can re-execute this query. - cache.complete(key, result, dep_node_index); + cache.complete(key, key_hash, result, dep_node_index); let job = { - let mut lock = state.active.lock_shard_by_value(&key); - match lock.remove(&key).unwrap() { - QueryResult::Started(job) => job, - QueryResult::Poisoned => panic!(), + let mut lock = state.active.lock_shard_by_hash(key_hash); + + match lock.raw_entry_mut().from_key_hashed_nocheck(key_hash, &key) { + RawEntryMut::Vacant(_) => panic!(), + RawEntryMut::Occupied(occupied) => match occupied.remove() { + QueryResult::Started(job) => job, + QueryResult::Poisoned => panic!(), + }, } }; @@ -211,7 +216,8 @@ where C: QueryCache, Tcx: DepContext, { - match cache.lookup(&key) { + let key_hash = sharded::make_hash(key); + match cache.lookup(&key, key_hash) { Some((value, index)) => { tcx.profiler().query_cache_hit(index.into()); tcx.dep_graph().read_index(index); @@ -248,6 +254,7 @@ fn wait_for_query( qcx: Qcx, span: Span, key: Q::Key, + key_hash: u64, latch: QueryLatch, current: Option, ) -> (Q::Value, Option) @@ -266,7 +273,7 @@ where match result { Ok(()) => { - let Some((v, index)) = query.query_cache(qcx).lookup(&key) else { + let Some((v, index)) = query.query_cache(qcx).lookup(&key, key_hash) else { cold_path(|| { // We didn't find the query result in the query cache. Check if it was // poisoned due to a panic instead. @@ -303,7 +310,8 @@ where Qcx: QueryContext, { let state = query.query_state(qcx); - let mut state_lock = state.active.lock_shard_by_value(&key); + let key_hash = sharded::make_hash(&key); + let mut state_lock = state.active.lock_shard_by_hash(key_hash); // For the parallel compiler we need to check both the query cache and query state structures // while holding the state lock to ensure that 1) the query has not yet completed and 2) the @@ -312,7 +320,7 @@ where // executing, but another thread may have already completed the query and stores it result // in the query cache. if cfg!(parallel_compiler) && qcx.dep_context().sess().threads() > 1 { - if let Some((value, index)) = query.query_cache(qcx).lookup(&key) { + if let Some((value, index)) = query.query_cache(qcx).lookup(&key, key_hash) { qcx.dep_context().profiler().query_cache_hit(index.into()); return (value, Some(index)); } @@ -320,20 +328,20 @@ where let current_job_id = qcx.current_query_job(); - match state_lock.entry(key) { - Entry::Vacant(entry) => { + match state_lock.raw_entry_mut().from_key_hashed_nocheck(key_hash, &key) { + RawEntryMut::Vacant(entry) => { // Nothing has computed or is computing the query, so we start a new job and insert it in the // state map. let id = qcx.next_job_id(); let job = QueryJob::new(id, span, current_job_id); - entry.insert(QueryResult::Started(job)); + entry.insert_hashed_nocheck(key_hash, key, QueryResult::Started(job)); // Drop the lock before we start executing the query drop(state_lock); - execute_job::<_, _, INCR>(query, qcx, state, key, id, dep_node) + execute_job::<_, _, INCR>(query, qcx, state, key, key_hash, id, dep_node) } - Entry::Occupied(mut entry) => { + RawEntryMut::Occupied(mut entry) => { match entry.get_mut() { QueryResult::Started(job) => { #[cfg(parallel_compiler)] @@ -344,7 +352,15 @@ where // Only call `wait_for_query` if we're using a Rayon thread pool // as it will attempt to mark the worker thread as blocked. - return wait_for_query(query, qcx, span, key, latch, current_job_id); + return wait_for_query( + query, + qcx, + span, + key, + key_hash, + latch, + current_job_id, + ); } let id = job.id; @@ -366,6 +382,7 @@ fn execute_job( qcx: Qcx, state: &QueryState, key: Q::Key, + key_hash: u64, id: QueryJobId, dep_node: Option>, ) -> (Q::Value, Option) @@ -397,7 +414,7 @@ where // This can't happen, as query feeding adds the very dependencies to the fed query // as its feeding query had. So if the fed query is red, so is its feeder, which will // get evaluated first, and re-feed the query. - if let Some((cached_result, _)) = cache.lookup(&key) { + if let Some((cached_result, _)) = cache.lookup(&key, key_hash) { let Some(hasher) = query.hash_result() else { panic!( "no_hash fed query later has its value computed.\n\ @@ -429,7 +446,7 @@ where } } } - job_owner.complete(cache, result, dep_node_index); + job_owner.complete(cache, key_hash, result, dep_node_index); (result, Some(dep_node_index)) } @@ -832,7 +849,7 @@ pub fn force_query( { // We may be concurrently trying both execute and force a query. // Ensure that only one of them runs the query. - if let Some((_, index)) = query.query_cache(qcx).lookup(&key) { + if let Some((_, index)) = query.query_cache(qcx).lookup(&key, sharded::make_hash(&key)) { qcx.dep_context().profiler().query_cache_hit(index.into()); return; } From 936d8ad972fe0b9b64b07911e32efa228239dbf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Mon, 21 Aug 2023 22:45:55 +0200 Subject: [PATCH 2/2] Use `find_or_find_insert_slot` for query execution --- .../rustc_query_system/src/query/plumbing.rs | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs index 4268e57cf5644..e58040a6747ee 100644 --- a/compiler/rustc_query_system/src/query/plumbing.rs +++ b/compiler/rustc_query_system/src/query/plumbing.rs @@ -328,21 +328,38 @@ where let current_job_id = qcx.current_query_job(); - match state_lock.raw_entry_mut().from_key_hashed_nocheck(key_hash, &key) { - RawEntryMut::Vacant(entry) => { + match state_lock.raw_table_mut().find_or_find_insert_slot( + key_hash, + |(k, _)| *k == key, + |(k, _)| sharded::make_hash(k), + ) { + Err(free_slot) => { // Nothing has computed or is computing the query, so we start a new job and insert it in the // state map. let id = qcx.next_job_id(); let job = QueryJob::new(id, span, current_job_id); - entry.insert_hashed_nocheck(key_hash, key, QueryResult::Started(job)); + + // SAFETY: The slot is still valid as there's + // been no mutation to the table since we hold the lock. + unsafe { + state_lock.raw_table_mut().insert_in_slot( + key_hash, + free_slot, + (key, QueryResult::Started(job)), + ); + } // Drop the lock before we start executing the query drop(state_lock); execute_job::<_, _, INCR>(query, qcx, state, key, key_hash, id, dep_node) } - RawEntryMut::Occupied(mut entry) => { - match entry.get_mut() { + Ok(bucket) => { + // SAFETY: We know this bucket is still valid + // since we just got it from `find_or_find_insert_slot`. + let entry = unsafe { &mut bucket.as_mut().1 }; + + match entry { QueryResult::Started(job) => { #[cfg(parallel_compiler)] if sync::is_dyn_thread_safe() {