From 8fc160b742e53eeb77a14b3e247b06b778f3243d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Wed, 16 Aug 2023 13:50:31 +0200 Subject: [PATCH 1/3] Add optimized lock methods for `Sharded` --- compiler/rustc_data_structures/src/sharded.rs | 53 ++++++++++- .../rustc_data_structures/src/sync/lock.rs | 87 ++++++++++++++++--- .../rustc_query_system/src/dep_graph/graph.rs | 12 +-- .../rustc_query_system/src/query/caches.rs | 8 +- .../rustc_query_system/src/query/plumbing.rs | 6 +- 5 files changed, 133 insertions(+), 33 deletions(-) diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs index 0f769c1f3bf4e..3706f62383fef 100644 --- a/compiler/rustc_data_structures/src/sharded.rs +++ b/compiler/rustc_data_structures/src/sharded.rs @@ -73,6 +73,53 @@ impl Sharded { } } + /// The shard is selected by hashing `val` with `FxHasher`. + #[inline] + pub fn lock_shard_by_value(&self, _val: &K) -> LockGuard<'_, T> { + match self { + Self::Single(single) => { + // Syncronization is disabled so use the `lock_assume_no_sync` method optimized + // for that case. + + // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus + // `might_be_dyn_thread_safe` was also false. + unsafe { single.lock_assume_no_sync() } + } + #[cfg(parallel_compiler)] + Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)), + } + } + + #[inline] + pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> { + self.lock_shard_by_index(get_shard_hash(hash)) + } + + #[inline] + pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> { + match self { + Self::Single(single) => { + // Syncronization is disabled so use the `lock_assume_no_sync` method optimized + // for that case. + + // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus + // `might_be_dyn_thread_safe` was also false. + unsafe { single.lock_assume_no_sync() } + } + #[cfg(parallel_compiler)] + Self::Shards(shards) => { + // Syncronization is enabled so use the `lock_assume_sync` method optimized + // for that case. + + // SAFETY (get_unchecked): The index gets ANDed with the shard mask, ensuring it is + // always inbounds. + // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating + // the lock thus `might_be_dyn_thread_safe` was also true. + unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume_sync() } + } + } + } + #[inline] pub fn lock_shards(&self) -> impl Iterator> { match self { @@ -124,7 +171,7 @@ impl ShardedHashMap { Q: Hash + Eq, { let hash = make_hash(value); - let mut shard = self.get_shard_by_hash(hash).lock(); + let mut shard = self.lock_shard_by_hash(hash); let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, value); match entry { @@ -144,7 +191,7 @@ impl ShardedHashMap { Q: Hash + Eq, { let hash = make_hash(&value); - let mut shard = self.get_shard_by_hash(hash).lock(); + let mut shard = self.lock_shard_by_hash(hash); let entry = shard.raw_entry_mut().from_key_hashed_nocheck(hash, &value); match entry { @@ -166,7 +213,7 @@ pub trait IntoPointer { impl ShardedHashMap { pub fn contains_pointer_to(&self, value: &T) -> bool { let hash = make_hash(&value); - let shard = self.get_shard_by_hash(hash).lock(); + let shard = self.lock_shard_by_hash(hash); let value = value.into_pointer(); shard.raw_entry().from_hash(hash, |entry| entry.into_pointer() == value).is_some() } diff --git a/compiler/rustc_data_structures/src/sync/lock.rs b/compiler/rustc_data_structures/src/sync/lock.rs index 62cd1b993de29..0a8c0306fe543 100644 --- a/compiler/rustc_data_structures/src/sync/lock.rs +++ b/compiler/rustc_data_structures/src/sync/lock.rs @@ -49,6 +49,20 @@ impl Lock { self.0.try_borrow_mut().ok() } + #[inline(always)] + #[track_caller] + // This is unsafe to match the API for the `parallel_compiler` case. + pub unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { + self.0.borrow_mut() + } + + #[inline(always)] + #[track_caller] + // This is unsafe to match the API for the `parallel_compiler` case. + pub unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { + self.0.borrow_mut() + } + #[inline(always)] #[track_caller] pub fn lock(&self) -> LockGuard<'_, T> { @@ -150,24 +164,45 @@ impl LockRaw { #[inline(always)] fn lock(&self) { - if super::ERROR_CHECKING { - // We're in the debugging mode, so assert that the lock is not held so we - // get a panic instead of waiting for the lock. - assert_eq!(self.try_lock(), true, "lock must not be hold"); - } else { - // SAFETY: This is safe since the union fields are used in accordance with `self.sync`. - unsafe { - if likely(!self.sync) { - if unlikely(self.opt.cell.replace(true)) { - cold_path(|| panic!("lock was already held")) - } - } else { - self.opt.lock.lock(); - } + // SAFETY: This is safe since `self.sync` is used in accordance with the preconditions of + // `lock_assume_no_sync` and `lock_assume_sync`. + unsafe { + if likely(!self.sync) { + self.lock_assume_no_sync() + } else { + self.lock_assume_sync(); + } + } + } + + /// This acquires the lock assuming no syncronization is required. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. + #[inline(always)] + unsafe fn lock_assume_no_sync(&self) { + // SAFETY: This is safe since `self.opt.cell` is the union field used due to the + // precondition on this function. + unsafe { + if unlikely(self.opt.cell.replace(true)) { + cold_path(|| panic!("lock was already held")) } } } + /// This acquires the lock assuming syncronization is required. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. + #[inline(always)] + unsafe fn lock_assume_sync(&self) { + // SAFETY: This is safe since `self.opt.lock` is the union field used due to the + // precondition on this function. + unsafe { + self.opt.lock.lock(); + } + } + /// This unlocks the lock. /// /// Safety @@ -217,6 +252,30 @@ impl Lock { if self.raw.try_lock() { Some(LockGuard { lock: self, marker: PhantomData }) } else { None } } + /// This acquires the lock assuming no syncronization is required. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. + #[inline(always)] + pub(crate) unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { + unsafe { + self.raw.lock_assume_no_sync(); + } + LockGuard { lock: self, marker: PhantomData } + } + + /// This acquires the lock assuming syncronization is required. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. + #[inline(always)] + pub(crate) unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { + unsafe { + self.raw.lock_assume_sync(); + } + LockGuard { lock: self, marker: PhantomData } + } + #[inline(always)] pub fn lock(&self) -> LockGuard<'_, T> { self.raw.lock(); diff --git a/compiler/rustc_query_system/src/dep_graph/graph.rs b/compiler/rustc_query_system/src/dep_graph/graph.rs index 7b7981e1425d3..fa54e1a2e6a74 100644 --- a/compiler/rustc_query_system/src/dep_graph/graph.rs +++ b/compiler/rustc_query_system/src/dep_graph/graph.rs @@ -629,12 +629,7 @@ impl DepGraphData { if let Some(prev_index) = self.previous.node_to_index_opt(dep_node) { self.current.prev_index_to_index.lock()[prev_index] } else { - self.current - .new_node_to_index - .get_shard_by_value(dep_node) - .lock() - .get(dep_node) - .copied() + self.current.new_node_to_index.lock_shard_by_value(dep_node).get(dep_node).copied() } } @@ -1201,8 +1196,7 @@ impl CurrentDepGraph { edges: EdgesVec, current_fingerprint: Fingerprint, ) -> DepNodeIndex { - let dep_node_index = match self.new_node_to_index.get_shard_by_value(&key).lock().entry(key) - { + let dep_node_index = match self.new_node_to_index.lock_shard_by_value(&key).entry(key) { Entry::Occupied(entry) => *entry.get(), Entry::Vacant(entry) => { let dep_node_index = @@ -1328,7 +1322,7 @@ impl CurrentDepGraph { ) { let node = &prev_graph.index_to_node(prev_index); debug_assert!( - !self.new_node_to_index.get_shard_by_value(node).lock().contains_key(node), + !self.new_node_to_index.lock_shard_by_value(node).contains_key(node), "node from previous graph present in new node collection" ); } diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs index d8aa377af4273..0240f012da053 100644 --- a/compiler/rustc_query_system/src/query/caches.rs +++ b/compiler/rustc_query_system/src/query/caches.rs @@ -55,7 +55,7 @@ where #[inline(always)] fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> { let key_hash = sharded::make_hash(key); - let lock = self.cache.get_shard_by_hash(key_hash).lock(); + 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 } @@ -63,7 +63,7 @@ where #[inline] fn complete(&self, key: K, value: V, index: DepNodeIndex) { - let mut lock = self.cache.get_shard_by_value(&key).lock(); + 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)); @@ -148,13 +148,13 @@ where #[inline(always)] fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> { - let lock = self.cache.get_shard_by_hash(key.index() as u64).lock(); + 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) { - let mut lock = self.cache.get_shard_by_hash(key.index() as u64).lock(); + 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 4fa168965a792..07db15e6d8be0 100644 --- a/compiler/rustc_query_system/src/query/plumbing.rs +++ b/compiler/rustc_query_system/src/query/plumbing.rs @@ -158,7 +158,7 @@ where cache.complete(key, result, dep_node_index); let job = { - let mut lock = state.active.get_shard_by_value(&key).lock(); + let mut lock = state.active.lock_shard_by_value(&key); match lock.remove(&key).unwrap() { QueryResult::Started(job) => job, QueryResult::Poisoned => panic!(), @@ -180,7 +180,7 @@ where // Poison the query so jobs waiting on it panic. let state = self.state; let job = { - let mut shard = state.active.get_shard_by_value(&self.key).lock(); + let mut shard = state.active.lock_shard_by_value(&self.key); let job = match shard.remove(&self.key).unwrap() { QueryResult::Started(job) => job, QueryResult::Poisoned => panic!(), @@ -303,7 +303,7 @@ where Qcx: QueryContext, { let state = query.query_state(qcx); - let mut state_lock = state.active.get_shard_by_value(&key).lock(); + let mut state_lock = state.active.lock_shard_by_value(&key); // 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 From 61cc00d23871b9d2e4c4f23af6917dc5ca7efb58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Sat, 2 Sep 2023 23:19:34 +0200 Subject: [PATCH 2/3] Refactor `Lock` implementation --- compiler/rustc_data_structures/src/sharded.rs | 11 +- compiler/rustc_data_structures/src/sync.rs | 3 +- .../rustc_data_structures/src/sync/lock.rs | 448 ++++++++---------- 3 files changed, 216 insertions(+), 246 deletions(-) diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs index 3706f62383fef..c56fce02af6d3 100644 --- a/compiler/rustc_data_structures/src/sharded.rs +++ b/compiler/rustc_data_structures/src/sharded.rs @@ -1,7 +1,7 @@ use crate::fx::{FxHashMap, FxHasher}; #[cfg(parallel_compiler)] use crate::sync::{is_dyn_thread_safe, CacheAligned}; -use crate::sync::{Lock, LockGuard}; +use crate::sync::{Assume, Lock, LockGuard}; #[cfg(parallel_compiler)] use itertools::Either; use std::borrow::Borrow; @@ -75,6 +75,7 @@ impl Sharded { /// The shard is selected by hashing `val` with `FxHasher`. #[inline] + #[track_caller] pub fn lock_shard_by_value(&self, _val: &K) -> LockGuard<'_, T> { match self { Self::Single(single) => { @@ -83,7 +84,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume_no_sync() } + unsafe { single.lock_assume(Assume::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)), @@ -91,11 +92,13 @@ impl Sharded { } #[inline] + #[track_caller] pub fn lock_shard_by_hash(&self, hash: u64) -> LockGuard<'_, T> { self.lock_shard_by_index(get_shard_hash(hash)) } #[inline] + #[track_caller] pub fn lock_shard_by_index(&self, _i: usize) -> LockGuard<'_, T> { match self { Self::Single(single) => { @@ -104,7 +107,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume_no_sync() } + unsafe { single.lock_assume(Assume::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(shards) => { @@ -115,7 +118,7 @@ impl Sharded { // always inbounds. // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating // the lock thus `might_be_dyn_thread_safe` was also true. - unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume_sync() } + unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Assume::Sync) } } } } diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs index 8eda9043e358b..a30ffcc120b01 100644 --- a/compiler/rustc_data_structures/src/sync.rs +++ b/compiler/rustc_data_structures/src/sync.rs @@ -49,7 +49,7 @@ use std::ops::{Deref, DerefMut}; use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe}; mod lock; -pub use lock::{Lock, LockGuard}; +pub use lock::{Assume, Lock, LockGuard}; mod worker_local; pub use worker_local::{Registry, WorkerLocal}; @@ -86,7 +86,6 @@ mod mode { // Whether thread safety might be enabled. #[inline] - #[cfg(parallel_compiler)] pub fn might_be_dyn_thread_safe() -> bool { DYN_THREAD_SAFE_MODE.load(Ordering::Relaxed) != DYN_NOT_THREAD_SAFE } diff --git a/compiler/rustc_data_structures/src/sync/lock.rs b/compiler/rustc_data_structures/src/sync/lock.rs index 0a8c0306fe543..0188eb114fbd7 100644 --- a/compiler/rustc_data_structures/src/sync/lock.rs +++ b/compiler/rustc_data_structures/src/sync/lock.rs @@ -3,283 +3,257 @@ //! //! When `cfg(parallel_compiler)` is not set, the lock is instead a wrapper around `RefCell`. -#[cfg(not(parallel_compiler))] -use std::cell::RefCell; -#[cfg(parallel_compiler)] -use { - crate::cold_path, - crate::sync::DynSend, - crate::sync::DynSync, - parking_lot::lock_api::RawMutex, - std::cell::Cell, - std::cell::UnsafeCell, - std::fmt, - std::intrinsics::{likely, unlikely}, - std::marker::PhantomData, - std::mem::ManuallyDrop, - std::ops::{Deref, DerefMut}, -}; +#![allow(dead_code)] -#[cfg(not(parallel_compiler))] -pub use std::cell::RefMut as LockGuard; +use std::fmt; #[cfg(not(parallel_compiler))] -#[derive(Debug)] -pub struct Lock(RefCell); - -#[cfg(not(parallel_compiler))] -impl Lock { - #[inline(always)] - pub fn new(inner: T) -> Self { - Lock(RefCell::new(inner)) - } +pub use disabled::*; +#[cfg(parallel_compiler)] +pub use enabled::*; - #[inline(always)] - pub fn into_inner(self) -> T { - self.0.into_inner() - } +#[derive(Clone, Copy, PartialEq)] +pub enum Assume { + NoSync, + Sync, +} - #[inline(always)] - pub fn get_mut(&mut self) -> &mut T { - self.0.get_mut() +mod enabled { + use super::Assume; + use crate::sync::mode; + #[cfg(parallel_compiler)] + use crate::sync::{DynSend, DynSync}; + use parking_lot::lock_api::RawMutex as _; + use parking_lot::RawMutex; + use std::cell::Cell; + use std::cell::UnsafeCell; + use std::hint::unreachable_unchecked; + use std::intrinsics::unlikely; + use std::marker::PhantomData; + use std::ops::{Deref, DerefMut}; + + /// A guard holding mutable access to a `Lock` which is in a locked state. + #[must_use = "if unused the Lock will immediately unlock"] + pub struct LockGuard<'a, T> { + lock: &'a Lock, + marker: PhantomData<&'a mut T>, + + /// The syncronization mode of the lock. This is explicitly passed to let LLVM relate it + /// to the original lock operation. + assume: Assume, } - #[inline(always)] - pub fn try_lock(&self) -> Option> { - self.0.try_borrow_mut().ok() + impl<'a, T: 'a> Deref for LockGuard<'a, T> { + type Target = T; + #[inline] + fn deref(&self) -> &T { + // SAFETY: We have shared access to the mutable access owned by this type, + // so we can give out a shared reference. + unsafe { &*self.lock.data.get() } + } } - #[inline(always)] - #[track_caller] - // This is unsafe to match the API for the `parallel_compiler` case. - pub unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() + impl<'a, T: 'a> DerefMut for LockGuard<'a, T> { + #[inline] + fn deref_mut(&mut self) -> &mut T { + // SAFETY: We have mutable access to the data so we can give out a mutable reference. + unsafe { &mut *self.lock.data.get() } + } } - #[inline(always)] - #[track_caller] - // This is unsafe to match the API for the `parallel_compiler` case. - pub unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() + impl<'a, T: 'a> Drop for LockGuard<'a, T> { + #[inline] + fn drop(&mut self) { + // SAFETY (dispatch): We get `self.assume` from the lock operation so it is consistent + // with the lock state. + // SAFETY (unlock): We know that the lock is locked as this type is a proof of that. + unsafe { + self.lock.dispatch( + self.assume, + |cell| { + debug_assert_eq!(cell.get(), true); + cell.set(false); + Some(()) + }, + |lock| lock.unlock(), + ); + }; + } } - #[inline(always)] - #[track_caller] - pub fn lock(&self) -> LockGuard<'_, T> { - self.0.borrow_mut() + enum LockMode { + NoSync(Cell), + Sync(RawMutex), } -} - -/// A guard holding mutable access to a `Lock` which is in a locked state. -#[cfg(parallel_compiler)] -#[must_use = "if unused the Lock will immediately unlock"] -pub struct LockGuard<'a, T> { - lock: &'a Lock, - marker: PhantomData<&'a mut T>, -} -#[cfg(parallel_compiler)] -impl<'a, T: 'a> Deref for LockGuard<'a, T> { - type Target = T; - #[inline] - fn deref(&self) -> &T { - // SAFETY: We have shared access to the mutable access owned by this type, - // so we can give out a shared reference. - unsafe { &*self.lock.data.get() } + impl LockMode { + #[inline(always)] + fn to_assume(&self) -> Assume { + match self { + LockMode::NoSync(..) => Assume::NoSync, + LockMode::Sync(..) => Assume::Sync, + } + } } -} -#[cfg(parallel_compiler)] -impl<'a, T: 'a> DerefMut for LockGuard<'a, T> { - #[inline] - fn deref_mut(&mut self) -> &mut T { - // SAFETY: We have mutable access to the data so we can give out a mutable reference. - unsafe { &mut *self.lock.data.get() } - } -} + /// The value representing a locked state for the `Cell`. + const LOCKED: bool = true; -#[cfg(parallel_compiler)] -impl<'a, T: 'a> Drop for LockGuard<'a, T> { - #[inline] - fn drop(&mut self) { - // SAFETY: We know that the lock is in a locked - // state because it is a invariant of this type. - unsafe { self.lock.raw.unlock() }; + /// A lock which only uses synchronization if `might_be_dyn_thread_safe` is true. + /// It implements `DynSend` and `DynSync` instead of the typical `Send` and `Sync`. + pub struct Lock { + mode: LockMode, + data: UnsafeCell, } -} -#[cfg(parallel_compiler)] -union LockRawUnion { - /// Indicates if the cell is locked. Only used if `LockRaw.sync` is false. - cell: ManuallyDrop>, + impl Lock { + #[inline(always)] + pub fn new(inner: T) -> Self { + Lock { + mode: if unlikely(mode::might_be_dyn_thread_safe()) { + // Create the lock with synchronization enabled using the `RawMutex` type. + LockMode::Sync(RawMutex::INIT) + } else { + // Create the lock with synchronization disabled. + LockMode::NoSync(Cell::new(!LOCKED)) + }, + data: UnsafeCell::new(inner), + } + } - /// A lock implementation that's only used if `LockRaw.sync` is true. - lock: ManuallyDrop, -} + #[inline(always)] + pub fn into_inner(self) -> T { + self.data.into_inner() + } -/// A raw lock which only uses synchronization if `might_be_dyn_thread_safe` is true. -/// It contains no associated data and is used in the implementation of `Lock` which does have such data. -/// -/// A manual implementation of a tagged union is used with the `sync` field and the `LockRawUnion` instead -/// of using enums as it results in better code generation. -#[cfg(parallel_compiler)] -struct LockRaw { - /// Indicates if synchronization is used via `opt.lock` if true, - /// or if a non-thread safe cell is used via `opt.cell`. This is set on initialization and never changed. - sync: bool, - opt: LockRawUnion, -} + #[inline(always)] + pub fn get_mut(&mut self) -> &mut T { + self.data.get_mut() + } -#[cfg(parallel_compiler)] -impl LockRaw { - fn new() -> Self { - if unlikely(super::mode::might_be_dyn_thread_safe()) { - // Create the lock with synchronization enabled using the `RawMutex` type. - LockRaw { - sync: true, - opt: LockRawUnion { lock: ManuallyDrop::new(parking_lot::RawMutex::INIT) }, + /// This dispatches on the `LockMode` and gives access to its variants depending on + /// `assume`. If `no_sync` returns `None` this will panic. + /// + /// Safety + /// This method must only be called if `might_be_dyn_thread_safe` on lock creation matches + /// matches the `assume` argument. + #[inline(always)] + #[track_caller] + unsafe fn dispatch( + &self, + assume: Assume, + no_sync: impl FnOnce(&Cell) -> Option, + sync: impl FnOnce(&RawMutex) -> R, + ) -> R { + #[inline(never)] + #[track_caller] + #[cold] + fn lock_held() -> ! { + panic!("lock was already held") } - } else { - // Create the lock with synchronization disabled. - LockRaw { sync: false, opt: LockRawUnion { cell: ManuallyDrop::new(Cell::new(false)) } } - } - } - #[inline(always)] - fn try_lock(&self) -> bool { - // SAFETY: This is safe since the union fields are used in accordance with `self.sync`. - unsafe { - if likely(!self.sync) { - if self.opt.cell.get() { - false - } else { - self.opt.cell.set(true); - true + match assume { + Assume::NoSync => { + let LockMode::NoSync(cell) = &self.mode else { + unsafe { unreachable_unchecked() } + }; + if let Some(v) = no_sync(cell) { + v + } else { + // Call this here instead of in `no_sync` so `track_caller` gets properly + // passed along. + lock_held() + } + } + Assume::Sync => { + let LockMode::Sync(lock) = &self.mode else { + unsafe { unreachable_unchecked() } + }; + sync(lock) } - } else { - self.opt.lock.try_lock() } } - } - #[inline(always)] - fn lock(&self) { - // SAFETY: This is safe since `self.sync` is used in accordance with the preconditions of - // `lock_assume_no_sync` and `lock_assume_sync`. - unsafe { - if likely(!self.sync) { - self.lock_assume_no_sync() - } else { - self.lock_assume_sync(); + #[inline(always)] + pub fn try_lock(&self) -> Option> { + let assume = self.mode.to_assume(); + unsafe { + self.dispatch( + assume, + |cell| Some((cell.get() != LOCKED).then(|| cell.set(LOCKED)).is_some()), + RawMutex::try_lock, + ) + .then(|| LockGuard { lock: self, marker: PhantomData, assume }) } } - } - /// This acquires the lock assuming no syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. - #[inline(always)] - unsafe fn lock_assume_no_sync(&self) { - // SAFETY: This is safe since `self.opt.cell` is the union field used due to the - // precondition on this function. - unsafe { - if unlikely(self.opt.cell.replace(true)) { - cold_path(|| panic!("lock was already held")) + #[inline(always)] + #[track_caller] + pub unsafe fn lock_assume(&self, assume: Assume) -> LockGuard<'_, T> { + unsafe { + self.dispatch( + assume, + |cell| (cell.replace(LOCKED) != LOCKED).then(|| ()), + RawMutex::lock, + ); + LockGuard { lock: self, marker: PhantomData, assume } } } - } - /// This acquires the lock assuming syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. - #[inline(always)] - unsafe fn lock_assume_sync(&self) { - // SAFETY: This is safe since `self.opt.lock` is the union field used due to the - // precondition on this function. - unsafe { - self.opt.lock.lock(); + #[inline(always)] + #[track_caller] + pub fn lock(&self) -> LockGuard<'_, T> { + unsafe { self.lock_assume(self.mode.to_assume()) } } } - /// This unlocks the lock. - /// - /// Safety - /// This method may only be called if the lock is currently held. - #[inline(always)] - unsafe fn unlock(&self) { - // SAFETY: The union use is safe since the union fields are used in accordance with - // `self.sync` and the `unlock` method precondition is upheld by the caller. - unsafe { - if likely(!self.sync) { - debug_assert_eq!(self.opt.cell.get(), true); - self.opt.cell.set(false); - } else { - self.opt.lock.unlock(); - } - } - } + #[cfg(parallel_compiler)] + unsafe impl DynSend for Lock {} + #[cfg(parallel_compiler)] + unsafe impl DynSync for Lock {} } -/// A lock which only uses synchronization if `might_be_dyn_thread_safe` is true. -/// It implements `DynSend` and `DynSync` instead of the typical `Send` and `Sync`. -#[cfg(parallel_compiler)] -pub struct Lock { - raw: LockRaw, - data: UnsafeCell, -} +mod disabled { + use super::Assume; + use std::cell::RefCell; -#[cfg(parallel_compiler)] -impl Lock { - #[inline(always)] - pub fn new(inner: T) -> Self { - Lock { raw: LockRaw::new(), data: UnsafeCell::new(inner) } - } + pub use std::cell::RefMut as LockGuard; - #[inline(always)] - pub fn into_inner(self) -> T { - self.data.into_inner() - } + pub struct Lock(RefCell); - #[inline(always)] - pub fn get_mut(&mut self) -> &mut T { - self.data.get_mut() - } + impl Lock { + #[inline(always)] + pub fn new(inner: T) -> Self { + Lock(RefCell::new(inner)) + } - #[inline(always)] - pub fn try_lock(&self) -> Option> { - if self.raw.try_lock() { Some(LockGuard { lock: self, marker: PhantomData }) } else { None } - } + #[inline(always)] + pub fn into_inner(self) -> T { + self.0.into_inner() + } - /// This acquires the lock assuming no syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was false on lock creation. - #[inline(always)] - pub(crate) unsafe fn lock_assume_no_sync(&self) -> LockGuard<'_, T> { - unsafe { - self.raw.lock_assume_no_sync(); + #[inline(always)] + pub fn get_mut(&mut self) -> &mut T { + self.0.get_mut() } - LockGuard { lock: self, marker: PhantomData } - } - /// This acquires the lock assuming syncronization is required. - /// - /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` was true on lock creation. - #[inline(always)] - pub(crate) unsafe fn lock_assume_sync(&self) -> LockGuard<'_, T> { - unsafe { - self.raw.lock_assume_sync(); + #[inline(always)] + pub fn try_lock(&self) -> Option> { + self.0.try_borrow_mut().ok() } - LockGuard { lock: self, marker: PhantomData } - } - #[inline(always)] - pub fn lock(&self) -> LockGuard<'_, T> { - self.raw.lock(); - LockGuard { lock: self, marker: PhantomData } + #[inline(always)] + #[track_caller] + // This is unsafe to match the API for the `parallel_compiler` case. + pub unsafe fn lock_assume(&self, _assume: Assume) -> LockGuard<'_, T> { + self.0.borrow_mut() + } + + #[inline(always)] + #[track_caller] + pub fn lock(&self) -> LockGuard<'_, T> { + self.0.borrow_mut() + } } } @@ -303,12 +277,13 @@ impl Lock { } } -#[cfg(parallel_compiler)] -unsafe impl DynSend for Lock {} -#[cfg(parallel_compiler)] -unsafe impl DynSync for Lock {} +impl Default for Lock { + #[inline] + fn default() -> Self { + Lock::new(T::default()) + } +} -#[cfg(parallel_compiler)] impl fmt::Debug for Lock { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.try_lock() { @@ -326,10 +301,3 @@ impl fmt::Debug for Lock { } } } - -impl Default for Lock { - #[inline] - fn default() -> Self { - Lock::new(T::default()) - } -} From 9690142bef648b88a4f8d87000f5edeacb0cc0a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= Date: Fri, 8 Sep 2023 09:27:25 +0200 Subject: [PATCH 3/3] Remove the `LockMode` enum and `dispatch` --- compiler/rustc_data_structures/src/sharded.rs | 8 +- compiler/rustc_data_structures/src/sync.rs | 2 +- .../rustc_data_structures/src/sync/lock.rs | 176 ++++++++---------- 3 files changed, 79 insertions(+), 107 deletions(-) diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs index c56fce02af6d3..29516fffd6a62 100644 --- a/compiler/rustc_data_structures/src/sharded.rs +++ b/compiler/rustc_data_structures/src/sharded.rs @@ -1,7 +1,7 @@ use crate::fx::{FxHashMap, FxHasher}; #[cfg(parallel_compiler)] use crate::sync::{is_dyn_thread_safe, CacheAligned}; -use crate::sync::{Assume, Lock, LockGuard}; +use crate::sync::{Lock, LockGuard, Mode}; #[cfg(parallel_compiler)] use itertools::Either; use std::borrow::Borrow; @@ -84,7 +84,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume(Assume::NoSync) } + unsafe { single.lock_assume(Mode::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(..) => self.lock_shard_by_hash(make_hash(_val)), @@ -107,7 +107,7 @@ impl Sharded { // SAFETY: We know `is_dyn_thread_safe` was false when creating the lock thus // `might_be_dyn_thread_safe` was also false. - unsafe { single.lock_assume(Assume::NoSync) } + unsafe { single.lock_assume(Mode::NoSync) } } #[cfg(parallel_compiler)] Self::Shards(shards) => { @@ -118,7 +118,7 @@ impl Sharded { // always inbounds. // SAFETY (lock_assume_sync): We know `is_dyn_thread_safe` was true when creating // the lock thus `might_be_dyn_thread_safe` was also true. - unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Assume::Sync) } + unsafe { shards.get_unchecked(_i & (SHARDS - 1)).0.lock_assume(Mode::Sync) } } } } diff --git a/compiler/rustc_data_structures/src/sync.rs b/compiler/rustc_data_structures/src/sync.rs index a30ffcc120b01..63f137516bd3b 100644 --- a/compiler/rustc_data_structures/src/sync.rs +++ b/compiler/rustc_data_structures/src/sync.rs @@ -49,7 +49,7 @@ use std::ops::{Deref, DerefMut}; use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe}; mod lock; -pub use lock::{Assume, Lock, LockGuard}; +pub use lock::{Lock, LockGuard, Mode}; mod worker_local; pub use worker_local::{Registry, WorkerLocal}; diff --git a/compiler/rustc_data_structures/src/sync/lock.rs b/compiler/rustc_data_structures/src/sync/lock.rs index 0188eb114fbd7..339aebbf81aae 100644 --- a/compiler/rustc_data_structures/src/sync/lock.rs +++ b/compiler/rustc_data_structures/src/sync/lock.rs @@ -7,19 +7,19 @@ use std::fmt; -#[cfg(not(parallel_compiler))] -pub use disabled::*; #[cfg(parallel_compiler)] -pub use enabled::*; +pub use maybe_sync::*; +#[cfg(not(parallel_compiler))] +pub use no_sync::*; #[derive(Clone, Copy, PartialEq)] -pub enum Assume { +pub enum Mode { NoSync, Sync, } -mod enabled { - use super::Assume; +mod maybe_sync { + use super::Mode; use crate::sync::mode; #[cfg(parallel_compiler)] use crate::sync::{DynSend, DynSync}; @@ -27,9 +27,9 @@ mod enabled { use parking_lot::RawMutex; use std::cell::Cell; use std::cell::UnsafeCell; - use std::hint::unreachable_unchecked; use std::intrinsics::unlikely; use std::marker::PhantomData; + use std::mem::ManuallyDrop; use std::ops::{Deref, DerefMut}; /// A guard holding mutable access to a `Lock` which is in a locked state. @@ -40,7 +40,7 @@ mod enabled { /// The syncronization mode of the lock. This is explicitly passed to let LLVM relate it /// to the original lock operation. - assume: Assume, + mode: Mode, } impl<'a, T: 'a> Deref for LockGuard<'a, T> { @@ -64,36 +64,26 @@ mod enabled { impl<'a, T: 'a> Drop for LockGuard<'a, T> { #[inline] fn drop(&mut self) { - // SAFETY (dispatch): We get `self.assume` from the lock operation so it is consistent - // with the lock state. - // SAFETY (unlock): We know that the lock is locked as this type is a proof of that. - unsafe { - self.lock.dispatch( - self.assume, - |cell| { - debug_assert_eq!(cell.get(), true); - cell.set(false); - Some(()) - }, - |lock| lock.unlock(), - ); - }; + // SAFETY (union access): We get `self.mode` from the lock operation so it is consistent + // with the `lock.mode` state. This means we access the right union fields. + match self.mode { + Mode::NoSync => { + let cell = unsafe { &self.lock.mode_union.no_sync }; + debug_assert_eq!(cell.get(), true); + cell.set(false); + } + // SAFETY (unlock): We know that the lock is locked as this type is a proof of that. + Mode::Sync => unsafe { self.lock.mode_union.sync.unlock() }, + } } } - enum LockMode { - NoSync(Cell), - Sync(RawMutex), - } + union ModeUnion { + /// Indicates if the cell is locked. Only used if `Lock.mode` is `NoSync`. + no_sync: ManuallyDrop>, - impl LockMode { - #[inline(always)] - fn to_assume(&self) -> Assume { - match self { - LockMode::NoSync(..) => Assume::NoSync, - LockMode::Sync(..) => Assume::Sync, - } - } + /// A lock implementation that's only used if `Lock.mode` is `Sync`. + sync: ManuallyDrop, } /// The value representing a locked state for the `Cell`. @@ -102,23 +92,26 @@ mod enabled { /// A lock which only uses synchronization if `might_be_dyn_thread_safe` is true. /// It implements `DynSend` and `DynSync` instead of the typical `Send` and `Sync`. pub struct Lock { - mode: LockMode, + /// Indicates if synchronization is used via `mode_union.sync` if it's `Sync`, or if a + /// not thread safe cell is used via `mode_union.no_sync` if it's `NoSync`. + /// This is set on initialization and never changed. + mode: Mode, + + mode_union: ModeUnion, data: UnsafeCell, } impl Lock { #[inline(always)] pub fn new(inner: T) -> Self { - Lock { - mode: if unlikely(mode::might_be_dyn_thread_safe()) { - // Create the lock with synchronization enabled using the `RawMutex` type. - LockMode::Sync(RawMutex::INIT) - } else { - // Create the lock with synchronization disabled. - LockMode::NoSync(Cell::new(!LOCKED)) - }, - data: UnsafeCell::new(inner), - } + let (mode, mode_union) = if unlikely(mode::might_be_dyn_thread_safe()) { + // Create the lock with synchronization enabled using the `RawMutex` type. + (Mode::Sync, ModeUnion { sync: ManuallyDrop::new(RawMutex::INIT) }) + } else { + // Create the lock with synchronization disabled. + (Mode::NoSync, ModeUnion { no_sync: ManuallyDrop::new(Cell::new(!LOCKED)) }) + }; + Lock { mode, mode_union, data: UnsafeCell::new(inner) } } #[inline(always)] @@ -131,20 +124,32 @@ mod enabled { self.data.get_mut() } - /// This dispatches on the `LockMode` and gives access to its variants depending on - /// `assume`. If `no_sync` returns `None` this will panic. + #[inline(always)] + pub fn try_lock(&self) -> Option> { + let mode = self.mode; + // SAFETY: This is safe since the union fields are used in accordance with `self.mode`. + match mode { + Mode::NoSync => { + let cell = unsafe { &self.mode_union.no_sync }; + let was_unlocked = cell.get() != LOCKED; + if was_unlocked { + cell.set(LOCKED); + } + was_unlocked + } + Mode::Sync => unsafe { self.mode_union.sync.try_lock() }, + } + .then(|| LockGuard { lock: self, marker: PhantomData, mode }) + } + + /// This acquires the lock assuming syncronization is in a specific mode. /// /// Safety - /// This method must only be called if `might_be_dyn_thread_safe` on lock creation matches - /// matches the `assume` argument. + /// This method must only be called with `Mode::Sync` if `might_be_dyn_thread_safe` was + /// true on lock creation. #[inline(always)] #[track_caller] - unsafe fn dispatch( - &self, - assume: Assume, - no_sync: impl FnOnce(&Cell) -> Option, - sync: impl FnOnce(&RawMutex) -> R, - ) -> R { + pub unsafe fn lock_assume(&self, mode: Mode) -> LockGuard<'_, T> { #[inline(never)] #[track_caller] #[cold] @@ -152,58 +157,25 @@ mod enabled { panic!("lock was already held") } - match assume { - Assume::NoSync => { - let LockMode::NoSync(cell) = &self.mode else { - unsafe { unreachable_unchecked() } - }; - if let Some(v) = no_sync(cell) { - v - } else { - // Call this here instead of in `no_sync` so `track_caller` gets properly - // passed along. - lock_held() + // SAFETY: This is safe since the union fields are used in accordance with `mode` + // which also must match `self.mode` due to the safety precondition. + unsafe { + match mode { + Mode::NoSync => { + if unlikely(self.mode_union.no_sync.replace(LOCKED) == LOCKED) { + lock_held() + } } + Mode::Sync => self.mode_union.sync.lock(), } - Assume::Sync => { - let LockMode::Sync(lock) = &self.mode else { - unsafe { unreachable_unchecked() } - }; - sync(lock) - } - } - } - - #[inline(always)] - pub fn try_lock(&self) -> Option> { - let assume = self.mode.to_assume(); - unsafe { - self.dispatch( - assume, - |cell| Some((cell.get() != LOCKED).then(|| cell.set(LOCKED)).is_some()), - RawMutex::try_lock, - ) - .then(|| LockGuard { lock: self, marker: PhantomData, assume }) - } - } - - #[inline(always)] - #[track_caller] - pub unsafe fn lock_assume(&self, assume: Assume) -> LockGuard<'_, T> { - unsafe { - self.dispatch( - assume, - |cell| (cell.replace(LOCKED) != LOCKED).then(|| ()), - RawMutex::lock, - ); - LockGuard { lock: self, marker: PhantomData, assume } } + LockGuard { lock: self, marker: PhantomData, mode } } #[inline(always)] #[track_caller] pub fn lock(&self) -> LockGuard<'_, T> { - unsafe { self.lock_assume(self.mode.to_assume()) } + unsafe { self.lock_assume(self.mode) } } } @@ -213,8 +185,8 @@ mod enabled { unsafe impl DynSync for Lock {} } -mod disabled { - use super::Assume; +mod no_sync { + use super::Mode; use std::cell::RefCell; pub use std::cell::RefMut as LockGuard; @@ -245,7 +217,7 @@ mod disabled { #[inline(always)] #[track_caller] // This is unsafe to match the API for the `parallel_compiler` case. - pub unsafe fn lock_assume(&self, _assume: Assume) -> LockGuard<'_, T> { + pub unsafe fn lock_assume(&self, _mode: Mode) -> LockGuard<'_, T> { self.0.borrow_mut() }