Store interened wrapped layout of lambda set on LambdaSet struct

The `LambdaSet` struct is frequently used independently to examine how a
lambda set should be packed or unpacked. However, it is also often
converted into a full layout via `Layout::LambdaSet(LambdaSet)` to be a
part of function arguments, for example.

In preparing to intern all layouts, we need a way to cheaply go from a
`lambda_set` to an interned `Layout::LambdaSet(lambda_set)`, since this
is a very common operation. The proposed solution is to keep the wrapped
layout cached on `LambdaSet` itself, which this PR does.

The tricky bit of inserting a lambda set is we need to fill in the
interned `full_layout` only after the lambda set is inserted,
but we don't want to allocate a new interned slot if the same lambda set
layout has already been inserted with a different `full_layout` slot.

For example, if we insert `LambdaSet { set : [A] }` twice in two
different threads, we want the `full_layout` they map to to be the same.
So we nede to check if an interned representation with a full_layout
exists, before we allocate a new full_layout and insert a fresh lambda
set.

So,
  - check if the "normalized" lambda set (with a void full_layout slot) maps to an
    inserted lambda set in
    - in a thread-local cache, or globally
  - if so, use that one immediately
  - otherwise, allocate a new (global) slot, intern the lambda set, and then fill the slot in
    - save the interned layout and lambda set mapping thread-locally
This commit is contained in:
Ayaz Hafiz 2023-01-03 16:02:22 -06:00
parent ce717dca8b
commit 26f08c999c
No known key found for this signature in database
GPG key ID: 0E2A37416A25EF58
3 changed files with 177 additions and 18 deletions

View file

@ -8,9 +8,10 @@ use std::{
use parking_lot::{Mutex, RwLock};
use roc_builtins::bitcode::{FloatWidth, IntWidth};
use roc_collections::{default_hasher, BumpMap};
use roc_module::symbol::Symbol;
use roc_target::TargetInfo;
use super::{Builtin, Layout};
use super::{Builtin, LambdaSet, Layout};
#[allow(unused)] // for now
pub struct InLayouts(PhantomData<()>);
@ -102,6 +103,14 @@ pub trait LayoutInterner<'a>: Sized {
// allocations when values already have interned representations.
fn insert(&mut self, value: Layout<'a>) -> InLayout<'a>;
/// Creates a [LambdaSet], including caching the [Layout::LambdaSet] representation of the
/// lambda set onto itself.
fn insert_lambda_set(
&mut self,
set: &'a &'a [(Symbol, &'a [InLayout<'a>])],
representation: InLayout<'a>,
) -> LambdaSet<'a>;
/// Retrieves a value from the interner.
fn get(&self, key: InLayout<'a>) -> Layout<'a>;
@ -159,6 +168,7 @@ pub struct GlobalLayoutInterner<'a>(Arc<GlobalLayoutInternerInner<'a>>);
#[derive(Debug)]
struct GlobalLayoutInternerInner<'a> {
map: Mutex<BumpMap<Layout<'a>, InLayout<'a>>>,
normalized_lambda_set_map: Mutex<BumpMap<LambdaSet<'a>, LambdaSet<'a>>>,
vec: RwLock<Vec<Layout<'a>>>,
}
@ -175,6 +185,7 @@ struct GlobalLayoutInternerInner<'a> {
pub struct TLLayoutInterner<'a> {
parent: GlobalLayoutInterner<'a>,
map: BumpMap<Layout<'a>, InLayout<'a>>,
normalized_lambda_set_map: BumpMap<LambdaSet<'a>, LambdaSet<'a>>,
/// Cache of interned values from the parent for local access.
vec: RefCell<Vec<Option<Layout<'a>>>>,
}
@ -186,6 +197,7 @@ pub struct TLLayoutInterner<'a> {
#[derive(Debug)]
pub struct STLayoutInterner<'a> {
map: BumpMap<Layout<'a>, InLayout<'a>>,
normalized_lambda_set_map: BumpMap<LambdaSet<'a>, LambdaSet<'a>>,
vec: Vec<Layout<'a>>,
}
@ -198,6 +210,18 @@ fn hash<V: std::hash::Hash>(val: V) -> u64 {
state.finish()
}
#[inline(always)]
fn make_normalized_lamdba_set<'a>(
set: &'a &'a [(Symbol, &'a [InLayout<'a>])],
representation: InLayout<'a>,
) -> LambdaSet<'a> {
LambdaSet {
set,
representation,
full_layout: InLayouts::VOID,
}
}
impl<'a> GlobalLayoutInterner<'a> {
/// Creates a new global interner with the given capacity.
pub fn with_capacity(cap: usize) -> Self {
@ -209,6 +233,7 @@ impl<'a> GlobalLayoutInterner<'a> {
TLLayoutInterner {
parent: Self(Arc::clone(&self.0)),
map: Default::default(),
normalized_lambda_set_map: Default::default(),
vec: Default::default(),
}
}
@ -217,13 +242,22 @@ impl<'a> GlobalLayoutInterner<'a> {
///
/// Returns an [Err] with `self` if there are outstanding references to the [GlobalLayoutInterner].
pub fn unwrap(self) -> Result<STLayoutInterner<'a>, Self> {
let GlobalLayoutInternerInner { map, vec } = match Arc::try_unwrap(self.0) {
let GlobalLayoutInternerInner {
map,
normalized_lambda_set_map,
vec,
} = match Arc::try_unwrap(self.0) {
Ok(inner) => inner,
Err(li) => return Err(Self(li)),
};
let map = Mutex::into_inner(map);
let normalized_lambda_set_map = Mutex::into_inner(normalized_lambda_set_map);
let vec = RwLock::into_inner(vec);
Ok(STLayoutInterner { map, vec })
Ok(STLayoutInterner {
map,
normalized_lambda_set_map,
vec,
})
}
/// Interns a value with a pre-computed hash.
@ -243,6 +277,35 @@ impl<'a> GlobalLayoutInterner<'a> {
*interned
}
fn get_or_insert_hashed_normalized_lambda_set(
&self,
normalized: LambdaSet<'a>,
normalized_hash: u64,
) -> WrittenGlobalLambdaSet<'a> {
let mut normalized_lambda_set_map = self.0.normalized_lambda_set_map.lock();
let (_, full_lambda_set) = normalized_lambda_set_map
.raw_entry_mut()
.from_key_hashed_nocheck(normalized_hash, &normalized)
.or_insert_with(|| {
// We don't already have an entry for the lambda set, which means it must be new to
// the world. Reserve a slot, insert the lambda set, and that should fill the slot
// in.
let mut vec = self.0.vec.write();
let slot = unsafe { InLayout::from_reserved_index(vec.len()) };
let lambda_set = LambdaSet {
full_layout: slot,
..normalized
};
vec.push(Layout::LambdaSet(lambda_set));
(normalized, lambda_set)
});
let full_layout = self.0.vec.read()[full_lambda_set.full_layout.0];
WrittenGlobalLambdaSet {
full_lambda_set: *full_lambda_set,
full_layout,
}
}
fn get(&self, interned: InLayout<'a>) -> Layout<'a> {
let InLayout(index, _) = interned;
self.0.vec.read()[index]
@ -253,6 +316,11 @@ impl<'a> GlobalLayoutInterner<'a> {
}
}
struct WrittenGlobalLambdaSet<'a> {
full_lambda_set: LambdaSet<'a>,
full_layout: Layout<'a>,
}
impl<'a> TLLayoutInterner<'a> {
/// Records an interned value in thread-specific storage, for faster access on lookups.
fn record(&self, key: Layout<'a>, interned: InLayout<'a>) {
@ -279,6 +347,51 @@ impl<'a> LayoutInterner<'a> for TLLayoutInterner<'a> {
interned
}
fn insert_lambda_set(
&mut self,
set: &'a &'a [(Symbol, &'a [InLayout<'a>])],
representation: InLayout<'a>,
) -> LambdaSet<'a> {
// The tricky bit of inserting a lambda set is we need to fill in the `full_layout` only
// after the lambda set is inserted, but we don't want to allocate a new interned slot if
// the same lambda set layout has already been inserted with a different `full_layout`
// slot.
//
// So,
// - check if the "normalized" lambda set (with a void full_layout slot) maps to an
// inserted lambda set in
// - in our thread-local cache, or globally
// - if so, use that one immediately
// - otherwise, allocate a new (global) slot, intern the lambda set, and then fill the slot in
let global = &self.parent;
let normalized = make_normalized_lamdba_set(set, representation);
let normalized_hash = hash(normalized);
let mut new_interned_layout = None;
let (_, &mut full_lambda_set) = self
.normalized_lambda_set_map
.raw_entry_mut()
.from_key_hashed_nocheck(normalized_hash, &normalized)
.or_insert_with(|| {
let WrittenGlobalLambdaSet {
full_lambda_set,
full_layout,
} = global.get_or_insert_hashed_normalized_lambda_set(normalized, normalized_hash);
// The Layout(lambda_set) isn't present in our thread; make sure it is for future
// reference.
new_interned_layout = Some((full_layout, full_lambda_set.full_layout));
(normalized, full_lambda_set)
});
if let Some((new_layout, new_interned)) = new_interned_layout {
// Write the interned lambda set layout into our thread-local cache.
self.record(new_layout, new_interned);
}
full_lambda_set
}
fn get(&self, key: InLayout<'a>) -> Layout<'a> {
if let Some(Some(value)) = self.vec.borrow().get(key.0) {
return *value;
@ -294,6 +407,7 @@ impl<'a> STLayoutInterner<'a> {
pub fn with_capacity(cap: usize) -> Self {
let mut interner = Self {
map: BumpMap::with_capacity_and_hasher(cap, default_hasher()),
normalized_lambda_set_map: BumpMap::with_capacity_and_hasher(cap, default_hasher()),
vec: Vec::with_capacity(cap),
};
fill_reserved_layouts(&mut interner);
@ -305,9 +419,14 @@ impl<'a> STLayoutInterner<'a> {
/// You should *only* use this if you need to go from a single-threaded to a concurrent context,
/// or in a case where you explicitly need access to [TLLayoutInterner]s.
pub fn into_global(self) -> GlobalLayoutInterner<'a> {
let STLayoutInterner { map, vec } = self;
let STLayoutInterner {
map,
normalized_lambda_set_map,
vec,
} = self;
GlobalLayoutInterner(Arc::new(GlobalLayoutInternerInner {
map: Mutex::new(map),
normalized_lambda_set_map: Mutex::new(normalized_lambda_set_map),
vec: RwLock::new(vec),
}))
}
@ -332,6 +451,37 @@ impl<'a> LayoutInterner<'a> for STLayoutInterner<'a> {
*interned
}
fn insert_lambda_set(
&mut self,
set: &'a &'a [(Symbol, &'a [InLayout<'a>])],
representation: InLayout<'a>,
) -> LambdaSet<'a> {
// IDEA:
// - check if the "normalized" lambda set (with a void full_layout slot) maps to an
// inserted lambda set
// - if so, use that one immediately
// - otherwise, allocate a new slot, intern the lambda set, and then fill the slot in
let normalized_lambda_set = make_normalized_lamdba_set(set, representation);
if let Some(lambda_set) = self.normalized_lambda_set_map.get(&normalized_lambda_set) {
return *lambda_set;
}
// This lambda set must be new to the interner, reserve a slot and fill it in.
let slot = unsafe { InLayout::from_reserved_index(self.vec.len()) };
let lambda_set = LambdaSet {
set,
representation,
full_layout: slot,
};
let filled_slot = self.insert(Layout::LambdaSet(lambda_set));
assert_eq!(slot, filled_slot);
self.normalized_lambda_set_map
.insert(normalized_lambda_set, lambda_set);
lambda_set
}
fn get(&self, key: InLayout<'a>) -> Layout<'a> {
let InLayout(index, _) = key;
self.vec[index]