Factor query arena allocation out from query caches

This commit is contained in:
John Kåre Alsaker
2023-02-08 19:53:48 +01:00
parent 4b34c7b766
commit a51a20531d
9 changed files with 160 additions and 254 deletions

View File

@@ -1,12 +1,10 @@
use crate::dep_graph::DepNodeIndex;
use rustc_arena::TypedArena;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::sharded;
#[cfg(parallel_compiler)]
use rustc_data_structures::sharded::Sharded;
use rustc_data_structures::sync::Lock;
use rustc_data_structures::sync::WorkerLocal;
use rustc_index::vec::{Idx, IndexVec};
use std::fmt::Debug;
use std::hash::Hash;
@@ -16,12 +14,10 @@ pub trait CacheSelector<'tcx, V> {
type Cache
where
V: Copy;
type ArenaCache;
}
pub trait QueryStorage {
type Value: Debug;
type Stored: Copy;
type Value: Copy;
}
pub trait QueryCache: QueryStorage + Sized {
@@ -31,9 +27,9 @@ pub trait QueryCache: QueryStorage + Sized {
/// It returns the shard index and a lock guard to the shard,
/// which will be used if the query is not in the cache and we need
/// to compute it.
fn lookup(&self, key: &Self::Key) -> Option<(Self::Stored, DepNodeIndex)>;
fn lookup(&self, key: &Self::Key) -> Option<(Self::Value, DepNodeIndex)>;
fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex) -> Self::Stored;
fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex);
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
}
@@ -44,7 +40,6 @@ impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<'tcx, V> for DefaultCacheSelecto
type Cache = DefaultCache<K, V>
where
V: Copy;
type ArenaCache = ArenaCache<'tcx, K, V>;
}
pub struct DefaultCache<K, V> {
@@ -62,7 +57,6 @@ impl<K, V> Default for DefaultCache<K, V> {
impl<K: Eq + Hash, V: Copy + Debug> QueryStorage for DefaultCache<K, V> {
type Value = V;
type Stored = V;
}
impl<K, V> QueryCache for DefaultCache<K, V>
@@ -85,7 +79,7 @@ where
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
fn complete(&self, key: K, value: V, index: DepNodeIndex) {
#[cfg(parallel_compiler)]
let mut lock = self.cache.get_shard_by_value(&key).lock();
#[cfg(not(parallel_compiler))]
@@ -93,7 +87,6 @@ where
// 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));
value
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
@@ -122,7 +115,6 @@ impl<'tcx, V: 'tcx> CacheSelector<'tcx, V> for SingleCacheSelector {
type Cache = SingleCache<V>
where
V: Copy;
type ArenaCache = ArenaCache<'tcx, (), V>;
}
pub struct SingleCache<V> {
@@ -137,7 +129,6 @@ impl<V> Default for SingleCache<V> {
impl<V: Copy + Debug> QueryStorage for SingleCache<V> {
type Value = V;
type Stored = V;
}
impl<V> QueryCache for SingleCache<V>
@@ -152,9 +143,8 @@ where
}
#[inline]
fn complete(&self, _key: (), value: V, index: DepNodeIndex) -> Self::Stored {
fn complete(&self, _key: (), value: V, index: DepNodeIndex) {
*self.cache.lock() = Some((value, index));
value
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
@@ -162,85 +152,12 @@ where
}
}
pub struct ArenaCache<'tcx, K, V> {
arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
#[cfg(parallel_compiler)]
cache: Sharded<FxHashMap<K, &'tcx (V, DepNodeIndex)>>,
#[cfg(not(parallel_compiler))]
cache: Lock<FxHashMap<K, &'tcx (V, DepNodeIndex)>>,
}
impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
fn default() -> Self {
ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), cache: Default::default() }
}
}
impl<'tcx, K: Eq + Hash, V: Debug + 'tcx> QueryStorage for ArenaCache<'tcx, K, V> {
type Value = V;
type Stored = &'tcx V;
}
impl<'tcx, K, V: 'tcx> QueryCache for ArenaCache<'tcx, K, V>
where
K: Eq + Hash + Clone + Debug,
V: Debug,
{
type Key = K;
#[inline(always)]
fn lookup(&self, key: &K) -> Option<(&'tcx V, DepNodeIndex)> {
let key_hash = sharded::make_hash(key);
#[cfg(parallel_compiler)]
let lock = self.cache.get_shard_by_hash(key_hash).lock();
#[cfg(not(parallel_compiler))]
let lock = self.cache.lock();
let result = lock.raw_entry().from_key_hashed_nocheck(key_hash, key);
if let Some((_, value)) = result { Some((&value.0, value.1)) } else { None }
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
let value = self.arena.alloc((value, index));
let value = unsafe { &*(value as *const _) };
#[cfg(parallel_compiler)]
let mut lock = self.cache.get_shard_by_value(&key).lock();
#[cfg(not(parallel_compiler))]
let mut lock = self.cache.lock();
// We may be overwriting another value. This is all right, since the dep-graph
// will check that the fingerprint matches.
lock.insert(key, value);
&value.0
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
#[cfg(parallel_compiler)]
{
let shards = self.cache.lock_shards();
for shard in shards.iter() {
for (k, v) in shard.iter() {
f(k, &v.0, v.1);
}
}
}
#[cfg(not(parallel_compiler))]
{
let map = self.cache.lock();
for (k, v) in map.iter() {
f(k, &v.0, v.1);
}
}
}
}
pub struct VecCacheSelector<K>(PhantomData<K>);
impl<'tcx, K: Idx, V: 'tcx> CacheSelector<'tcx, V> for VecCacheSelector<K> {
type Cache = VecCache<K, V>
where
V: Copy;
type ArenaCache = VecArenaCache<'tcx, K, V>;
}
pub struct VecCache<K: Idx, V> {
@@ -258,7 +175,6 @@ impl<K: Idx, V> Default for VecCache<K, V> {
impl<K: Eq + Idx, V: Copy + Debug> QueryStorage for VecCache<K, V> {
type Value = V;
type Stored = V;
}
impl<K, V> QueryCache for VecCache<K, V>
@@ -278,87 +194,12 @@ where
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
fn complete(&self, key: K, value: V, index: DepNodeIndex) {
#[cfg(parallel_compiler)]
let mut lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
#[cfg(not(parallel_compiler))]
let mut lock = self.cache.lock();
lock.insert(key, (value, index));
value
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
#[cfg(parallel_compiler)]
{
let shards = self.cache.lock_shards();
for shard in shards.iter() {
for (k, v) in shard.iter_enumerated() {
if let Some(v) = v {
f(&k, &v.0, v.1);
}
}
}
}
#[cfg(not(parallel_compiler))]
{
let map = self.cache.lock();
for (k, v) in map.iter_enumerated() {
if let Some(v) = v {
f(&k, &v.0, v.1);
}
}
}
}
}
pub struct VecArenaCache<'tcx, K: Idx, V> {
arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
#[cfg(parallel_compiler)]
cache: Sharded<IndexVec<K, Option<&'tcx (V, DepNodeIndex)>>>,
#[cfg(not(parallel_compiler))]
cache: Lock<IndexVec<K, Option<&'tcx (V, DepNodeIndex)>>>,
}
impl<'tcx, K: Idx, V> Default for VecArenaCache<'tcx, K, V> {
fn default() -> Self {
VecArenaCache {
arena: WorkerLocal::new(|_| TypedArena::default()),
cache: Default::default(),
}
}
}
impl<'tcx, K: Eq + Idx, V: Debug + 'tcx> QueryStorage for VecArenaCache<'tcx, K, V> {
type Value = V;
type Stored = &'tcx V;
}
impl<'tcx, K, V: 'tcx> QueryCache for VecArenaCache<'tcx, K, V>
where
K: Eq + Idx + Clone + Debug,
V: Debug,
{
type Key = K;
#[inline(always)]
fn lookup(&self, key: &K) -> Option<(&'tcx V, DepNodeIndex)> {
#[cfg(parallel_compiler)]
let lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
#[cfg(not(parallel_compiler))]
let lock = self.cache.lock();
if let Some(Some(value)) = lock.get(*key) { Some((&value.0, value.1)) } else { None }
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) -> Self::Stored {
let value = self.arena.alloc((value, index));
let value = unsafe { &*(value as *const _) };
#[cfg(parallel_compiler)]
let mut lock = self.cache.get_shard_by_hash(key.index() as u64).lock();
#[cfg(not(parallel_compiler))]
let mut lock = self.cache.lock();
lock.insert(key, value);
&value.0
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {