Skip to content

Commit

Permalink
Auto merge of rust-lang#115747 - Zoxc:query-hashes, r=<try>
Browse files Browse the repository at this point in the history
Optimize hash map operations in the query system

This optimizes hash map operations in the query system by explicitly passing hashes and using more optimal operations. `find_or_find_insert_slot` in particular saves a hash table lookup over `entry`. It's not yet available in a safe API, but will be in rust-lang/hashbrown#466.

<table><tr><td rowspan="2">Benchmark</td><td colspan="1"><b>Before</b></th><td colspan="2"><b>After</b></th></tr><tr><td align="right">Time</td><td align="right">Time</td><td align="right">%</th></tr><tr><td>🟣 <b>clap</b>:check</td><td align="right">1.6189s</td><td align="right">1.6129s</td><td align="right"> -0.37%</td></tr><tr><td>🟣 <b>hyper</b>:check</td><td align="right">0.2353s</td><td align="right">0.2337s</td><td align="right"> -0.67%</td></tr><tr><td>🟣 <b>regex</b>:check</td><td align="right">0.9344s</td><td align="right">0.9289s</td><td align="right"> -0.59%</td></tr><tr><td>🟣 <b>syn</b>:check</td><td align="right">1.4693s</td><td align="right">1.4652s</td><td align="right"> -0.28%</td></tr><tr><td>🟣 <b>syntex_syntax</b>:check</td><td align="right">5.6606s</td><td align="right">5.6439s</td><td align="right"> -0.30%</td></tr><tr><td>Total</td><td align="right">9.9185s</td><td align="right">9.8846s</td><td align="right"> -0.34%</td></tr><tr><td>Summary</td><td align="right">1.0000s</td><td align="right">0.9956s</td><td align="right"> -0.44%</td></tr></table>

r? `@cjgillot`
  • Loading branch information
bors committed Sep 21, 2023
2 parents 0fd7ce9 + 936d8ad commit d898dc6
Show file tree
Hide file tree
Showing 7 changed files with 91 additions and 41 deletions.
3 changes: 3 additions & 0 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -3600,6 +3600,7 @@ dependencies = [
"cfg-if",
"elsa",
"ena",
"hashbrown 0.14.0",
"indexmap 2.0.0",
"itertools",
"jobserver",
Expand Down Expand Up @@ -4328,7 +4329,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",
Expand Down
5 changes: 5 additions & 0 deletions compiler/rustc_data_structures/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
Expand Down
1 change: 1 addition & 0 deletions compiler/rustc_data_structures/src/marker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,7 @@ cfg_if!(
[std::sync::mpsc::Sender<T> where T: DynSend]
[std::sync::Arc<T> where T: ?Sized + DynSync + DynSend]
[std::sync::LazyLock<T, F> where T: DynSend, F: DynSend]
[hashbrown::HashMap<K, V, S> where K: DynSend, V: DynSend, S: DynSend]
[std::collections::HashSet<K, S> where K: DynSend, S: DynSend]
[std::collections::HashMap<K, V, S> where K: DynSend, V: DynSend, S: DynSend]
[std::collections::BTreeMap<K, V, A> where K: DynSend, V: DynSend, A: std::alloc::Allocator + Clone + DynSend]
Expand Down
3 changes: 2 additions & 1 deletion compiler/rustc_middle/src/query/plumbing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
Expand Down
6 changes: 6 additions & 0 deletions compiler/rustc_query_system/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
Expand Down
32 changes: 16 additions & 16 deletions compiler/rustc_query_system/src/query/caches.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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));
}
Expand All @@ -35,7 +35,7 @@ impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelecto
}

pub struct DefaultCache<K, V> {
cache: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
cache: Sharded<hashbrown::HashMap<K, (V, DepNodeIndex), BuildHasherDefault<FxHasher>>>,
}

impl<K, V> Default for DefaultCache<K, V> {
Expand All @@ -53,20 +53,20 @@ 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);

if let Some((_, value)) = result { Some(*value) } else { None }
}

#[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)) {
Expand Down Expand Up @@ -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();
}

Expand Down Expand Up @@ -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));
}
Expand Down
82 changes: 58 additions & 24 deletions compiler/rustc_query_system/src/query/plumbing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,27 @@ 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;

use super::QueryConfig;

pub struct QueryState<K, D: DepKind> {
active: Sharded<FxHashMap<K, QueryResult<D>>>,
active: Sharded<hashbrown::HashMap<K, QueryResult<D>, BuildHasherDefault<FxHasher>>>,
}

/// Indicates the state of a query for a given key in a query map.
Expand Down Expand Up @@ -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<C>(self, cache: &C, result: C::Value, dep_node_index: DepNodeIndex)
fn complete<C>(self, cache: &C, key_hash: u64, result: C::Value, dep_node_index: DepNodeIndex)
where
C: QueryCache<Key = K>,
{
Expand All @@ -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!(),
},
}
};

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -248,6 +254,7 @@ fn wait_for_query<Q, Qcx>(
qcx: Qcx,
span: Span,
key: Q::Key,
key_hash: u64,
latch: QueryLatch<Qcx::DepKind>,
current: Option<QueryJobId>,
) -> (Q::Value, Option<DepNodeIndex>)
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -312,29 +320,46 @@ 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));
}
}

let current_job_id = qcx.current_query_job();

match state_lock.entry(key) {
Entry::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(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, id, dep_node)
execute_job::<_, _, INCR>(query, qcx, state, key, key_hash, id, dep_node)
}
Entry::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() {
Expand All @@ -344,7 +369,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;
Expand All @@ -366,6 +399,7 @@ fn execute_job<Q, Qcx, const INCR: bool>(
qcx: Qcx,
state: &QueryState<Q::Key, Qcx::DepKind>,
key: Q::Key,
key_hash: u64,
id: QueryJobId,
dep_node: Option<DepNode<Qcx::DepKind>>,
) -> (Q::Value, Option<DepNodeIndex>)
Expand Down Expand Up @@ -397,7 +431,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\
Expand Down Expand Up @@ -429,7 +463,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))
}
Expand Down Expand Up @@ -832,7 +866,7 @@ pub fn force_query<Q, Qcx>(
{
// 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;
}
Expand Down

0 comments on commit d898dc6

Please sign in to comment.