Skip to content

Commit 9b72cc9

Browse files
committed
Auto merge of #115388 - Zoxc:sharded-lock, r=SparrowLii
Add optimized lock methods for `Sharded` and refactor `Lock` This adds methods to `Sharded` which pick a shard and also locks it. These branch on parallelism just once instead of twice, improving performance. Benchmark for `cfg(parallel_compiler)` and 1 thread: <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.6461s</td><td align="right">1.6345s</td><td align="right"> -0.70%</td></tr><tr><td>🟣 <b>hyper</b>:check</td><td align="right">0.2414s</td><td align="right">0.2394s</td><td align="right"> -0.83%</td></tr><tr><td>🟣 <b>regex</b>:check</td><td align="right">0.9205s</td><td align="right">0.9143s</td><td align="right"> -0.67%</td></tr><tr><td>🟣 <b>syn</b>:check</td><td align="right">1.4981s</td><td align="right">1.4869s</td><td align="right"> -0.75%</td></tr><tr><td>🟣 <b>syntex_syntax</b>:check</td><td align="right">5.7629s</td><td align="right">5.7256s</td><td align="right"> -0.65%</td></tr><tr><td>Total</td><td align="right">10.0690s</td><td align="right">10.0008s</td><td align="right"> -0.68%</td></tr><tr><td>Summary</td><td align="right">1.0000s</td><td align="right">0.9928s</td><td align="right"> -0.72%</td></tr></table> cc `@SparrowLii`
2 parents 9d311f9 + 9690142 commit 9b72cc9

File tree

6 files changed

+252
-210
lines changed

6 files changed

+252
-210
lines changed

compiler/rustc_data_structures/src/sharded.rs

+54-4
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::fx::{FxHashMap, FxHasher};
22
#[cfg(parallel_compiler)]
33
use crate::sync::{is_dyn_thread_safe, CacheAligned};
4-
use crate::sync::{Lock, LockGuard};
4+
use crate::sync::{Lock, LockGuard, Mode};
55
#[cfg(parallel_compiler)]
66
use itertools::Either;
77
use std::borrow::Borrow;
@@ -73,6 +73,56 @@ impl<T> Sharded<T> {
7373
}
7474
}
7575

76+
/// The shard is selected by hashing `val` with `FxHasher`.
77+
#[inline]
78+
#[track_caller]
79+
pub fn lock_shard_by_value<K: Hash + ?Sized>(&self, _val: &K) -> LockGuard<'_, T> {
80+
match self {
81+
Self::Single(single) => {
82+
// Syncronization is disabled so use the `lock_assume_no_sync` method optimized
83+
// for that case.
84+
85+
// SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
86+
// `might_be_dyn_thread_safe` was also false.
87+
unsafe { single.lock_assume(Mode::NoSync) }
88+
}
89+
#[cfg(parallel_compiler)]
90+
Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)),
91+
}
92+
}
93+
94+
#[inline]
95+
#[track_caller]
96+
pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> {
97+
self.lock_shard_by_index(get_shard_hash(hash))
98+
}
99+
100+
#[inline]
101+
#[track_caller]
102+
pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> {
103+
match self {
104+
Self::Single(single) => {
105+
// Syncronization is disabled so use the `lock_assume_no_sync` method optimized
106+
// for that case.
107+
108+
// SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus
109+
// `might_be_dyn_thread_safe` was also false.
110+
unsafe { single.lock_assume(Mode::NoSync) }
111+
}
112+
#[cfg(parallel_compiler)]
113+
Self::Shards(shards) => {
114+
// Syncronization is enabled so use the `lock_assume_sync` method optimized
115+
// for that case.
116+
117+
// SAFETY (get_unchecked): The index gets ANDed with the shard mask, ensuring it is
118+
// always inbounds.
119+
// SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating
120+
// the lock thus `might_be_dyn_thread_safe` was also true.
121+
unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Mode::Sync) }
122+
}
123+
}
124+
}
125+
76126
#[inline]
77127
pub fn lock_shards(&self) -> impl Iterator<Item = LockGuard<'_, T>> {
78128
match self {
@@ -124,7 +174,7 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
124174
Q: Hash + Eq,
125175
{
126176
let hash = make_hash(value);
127-
let mut shard = self.get_shard_by_hash(hash).lock();
177+
let mut shard = self.lock_shard_by_hash(hash);
128178
let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value);
129179

130180
match entry {
@@ -144,7 +194,7 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
144194
Q: Hash + Eq,
145195
{
146196
let hash = make_hash(&value);
147-
let mut shard = self.get_shard_by_hash(hash).lock();
197+
let mut shard = self.lock_shard_by_hash(hash);
148198
let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value);
149199

150200
match entry {
@@ -166,7 +216,7 @@ pub trait IntoPointer {
166216
impl<K: Eq + Hash + Copy + IntoPointer> ShardedHashMap<K, ()> {
167217
pub fn contains_pointer_to<T: Hash + IntoPointer>(&self, value: &T) -> bool {
168218
let hash = make_hash(&value);
169-
let shard = self.get_shard_by_hash(hash).lock();
219+
let shard = self.lock_shard_by_hash(hash);
170220
let value = value.into_pointer();
171221
shard.raw_entry().from_hash(hash, |entry| entry.into_pointer() == value).is_some()
172222
}

compiler/rustc_data_structures/src/sync.rs

+1-2
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ use std::ops::{Deref, DerefMut};
4949
use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
5050

5151
mod lock;
52-
pub use lock::{Lock, LockGuard};
52+
pub use lock::{Lock, LockGuard, Mode};
5353

5454
mod worker_local;
5555
pub use worker_local::{Registry, WorkerLocal};
@@ -86,7 +86,6 @@ mod mode {
8686

8787
// Whether thread safety might be enabled.
8888
#[inline]
89-
#[cfg(parallel_compiler)]
9089
pub fn might_be_dyn_thread_safe() -> bool {
9190
DYN_THREAD_SAFE_MODE.load(Ordering::Relaxed) != DYN_NOT_THREAD_SAFE
9291
}

0 commit comments

Comments
 (0)