Auto merge of #111755 - Zoxc:sharded-switch, r=cjgillot

Use only one shard with a single thread

This changes `Sharded` to only access a single shard using a mask set to `0` when a single thread is used, which leads to cache utilization improvements.

Performance improvement with 1 thread and `cfg(parallel_compiler)`:
<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.7402s</td><td align="right">1.7004s</td><td align="right">💚  -2.29%</td></tr><tr><td>🟣 <b>hyper</b>:check</td><td align="right">0.2633s</td><td align="right">0.2550s</td><td align="right">💚  -3.12%</td></tr><tr><td>🟣 <b>regex</b>:check</td><td align="right">0.9716s</td><td align="right">0.9482s</td><td align="right">💚  -2.41%</td></tr><tr><td>🟣 <b>syn</b>:check</td><td align="right">1.5679s</td><td align="right">1.5358s</td><td align="right">💚  -2.05%</td></tr><tr><td>🟣 <b>syntex_syntax</b>:check</td><td align="right">6.0569s</td><td align="right">5.9272s</td><td align="right">💚  -2.14%</td></tr><tr><td>Total</td><td align="right">10.5999s</td><td align="right">10.3666s</td><td align="right">💚  -2.20%</td></tr><tr><td>Summary</td><td align="right">1.0000s</td><td align="right">0.9760s</td><td align="right">💚  -2.40%</td></tr></table>

cc `@SparrowLii`

1 files changed, 38 insertions, 10 deletions

diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs
index 7ed70ba1e0f..40cbf14958e 100644
--- a/compiler/rustc_data_structures/src/sharded.rs
+++ b/compiler/rustc_data_structures/src/sharded.rs
@@ -1,4 +1,6 @@
 use crate::fx::{FxHashMap, FxHasher};
+#[cfg(parallel_compiler)]
+use crate::sync::is_dyn_thread_safe;
 use crate::sync::{CacheAligned, Lock, LockGuard};
 use std::borrow::Borrow;
 use std::collections::hash_map::RawEntryMut;
@@ -18,6 +20,11 @@ pub const SHARDS: usize = 1 << SHARD_BITS;
 
 /// An array of cache-line aligned inner locked structures with convenience methods.
 pub struct Sharded<T> {
+    /// This mask is used to ensure that accesses are inbounds of `shards`.
+    /// When dynamic thread safety is off, this field is set to 0 causing only
+    /// a single shard to be used for greater cache efficiency.
+    #[cfg(parallel_compiler)]
+    mask: usize,
     shards: [CacheAligned<Lock<T>>; SHARDS],
 }
 
@@ -31,31 +38,54 @@ impl<T: Default> Default for Sharded<T> {
 impl<T> Sharded<T> {
     #[inline]
     pub fn new(mut value: impl FnMut() -> T) -> Self {
-        Sharded { shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))) }
+        Sharded {
+            #[cfg(parallel_compiler)]
+            mask: if is_dyn_thread_safe() { SHARDS - 1 } else { 0 },
+            shards: [(); SHARDS].map(|()| CacheAligned(Lock::new(value()))),
+        }
+    }
+
+    #[inline(always)]
+    fn mask(&self) -> usize {
+        #[cfg(parallel_compiler)]
+        {
+            if SHARDS == 1 { 0 } else { self.mask }
+        }
+        #[cfg(not(parallel_compiler))]
+        {
+            0
+        }
+    }
+
+    #[inline(always)]
+    fn count(&self) -> usize {
+        // `self.mask` is always one below the used shard count
+        self.mask() + 1
     }
 
     /// The shard is selected by hashing `val` with `FxHasher`.
     #[inline]
     pub fn get_shard_by_value<K: Hash + ?Sized>(&self, val: &K) -> &Lock<T> {
-        if SHARDS == 1 { &self.shards[0].0 } else { self.get_shard_by_hash(make_hash(val)) }
+        self.get_shard_by_hash(if SHARDS == 1 { 0 } else { make_hash(val) })
     }
 
     #[inline]
     pub fn get_shard_by_hash(&self, hash: u64) -> &Lock<T> {
-        &self.shards[get_shard_index_by_hash(hash)].0
+        self.get_shard_by_index(get_shard_hash(hash))
     }
 
     #[inline]
     pub fn get_shard_by_index(&self, i: usize) -> &Lock<T> {
-        &self.shards[i].0
+        // SAFETY: The index get ANDed with the mask, ensuring it is always inbounds.
+        unsafe { &self.shards.get_unchecked(i & self.mask()).0 }
     }
 
     pub fn lock_shards(&self) -> Vec<LockGuard<'_, T>> {
-        (0..SHARDS).map(|i| self.shards[i].0.lock()).collect()
+        (0..self.count()).map(|i| self.get_shard_by_index(i).lock()).collect()
     }
 
     pub fn try_lock_shards(&self) -> Option<Vec<LockGuard<'_, T>>> {
-        (0..SHARDS).map(|i| self.shards[i].0.try_lock()).collect()
+        (0..self.count()).map(|i| self.get_shard_by_index(i).try_lock()).collect()
     }
 }
 
@@ -136,11 +166,9 @@ pub fn make_hash<K: Hash + ?Sized>(val: &K) -> u64 {
 /// `hash` can be computed with any hasher, so long as that hasher is used
 /// consistently for each `Sharded` instance.
 #[inline]
-#[allow(clippy::modulo_one)]
-pub fn get_shard_index_by_hash(hash: u64) -> usize {
+fn get_shard_hash(hash: u64) -> usize {
     let hash_len = mem::size_of::<usize>();
     // Ignore the top 7 bits as hashbrown uses these and get the next SHARD_BITS highest bits.
     // hashbrown also uses the lowest bits, so we can't use those
-    let bits = (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize;
-    bits % SHARDS
+    (hash >> (hash_len * 8 - 7 - SHARD_BITS)) as usize
 }