Skip to content

Commit f981e7d

Browse files
VeykrilChayimFriedman2
Veykril
and
ChayimFriedman2
authored
perf: Reduce memory usage by deduplicating type information (#803)
* Reduce memory usage by deduplicating type information We were storing the type information, 3 words wide, for each memo in each slot, while it is always constant wrt. the ingredient (different slots of the same ingredients will always have the same memos in the same order). This introduces some more unsafety, and the result wasn't as fast so I also had to use some lock-free structures, but the result is worth it: this shaves off 230mb from rust-analyzer with new Salsa. * Simplify * Replace `RwLock` with boxcar + `AtomicPtr` * Use TypeId and allocate instead * Use `OnceLock` instead of atomic-ptr --------- Co-authored-by: Chayim Refael Friedman <[email protected]>
1 parent cf9efae commit f981e7d

17 files changed

+472
-207
lines changed

‎components/salsa-macro-rules/src/setup_tracked_fn.rs

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -256,21 +256,45 @@ macro_rules! setup_tracked_fn {
256256
struct_index
257257
}
258258
};
259-
let memo_ingredient_indices = From::from((zalsa, struct_index, first_index));
260259

261-
let fn_ingredient = <$zalsa::function::IngredientImpl<$Configuration>>::new(
262-
first_index,
263-
memo_ingredient_indices,
264-
$lru,
265-
zalsa.views().downcaster_for::<dyn $Db>()
266-
);
260+
$zalsa::macro_if! { $needs_interner =>
261+
let intern_ingredient = <$zalsa::interned::IngredientImpl<$Configuration>>::new(
262+
first_index.successor(0)
263+
);
264+
}
265+
266+
let intern_ingredient_memo_types = $zalsa::macro_if! {
267+
if $needs_interner {
268+
Some($zalsa::Ingredient::memo_table_types(&intern_ingredient))
269+
} else {
270+
None
271+
}
272+
};
273+
// SAFETY: We call with the correct memo types.
274+
let memo_ingredient_indices = unsafe {
275+
$zalsa::NewMemoIngredientIndices::create(
276+
zalsa,
277+
struct_index,
278+
first_index,
279+
$zalsa::function::MemoEntryType::of::<$zalsa::function::Memo<$Configuration>>(),
280+
intern_ingredient_memo_types,
281+
)
282+
};
283+
284+
// SAFETY: We pass the MemoEntryType for this Configuration, and we lookup the memo types table correctly.
285+
let fn_ingredient = unsafe {
286+
<$zalsa::function::IngredientImpl<$Configuration>>::new(
287+
first_index,
288+
memo_ingredient_indices,
289+
$lru,
290+
zalsa.views().downcaster_for::<dyn $Db>(),
291+
)
292+
};
267293
$zalsa::macro_if! {
268294
if $needs_interner {
269295
vec![
270296
Box::new(fn_ingredient),
271-
Box::new(<$zalsa::interned::IngredientImpl<$Configuration>>::new(
272-
first_index.successor(0)
273-
)),
297+
Box::new(intern_ingredient),
274298
]
275299
} else {
276300
vec![

‎src/accumulator.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ use std::any::{Any, TypeId};
44
use std::fmt;
55
use std::marker::PhantomData;
66
use std::panic::UnwindSafe;
7+
use std::sync::Arc;
78

89
use accumulated::{Accumulated, AnyAccumulated};
910

1011
use crate::function::VerifyResult;
1112
use crate::ingredient::{fmt_index, Ingredient, Jar};
1213
use crate::plumbing::IngredientIndices;
14+
use crate::table::memo::MemoTableTypes;
1315
use crate::zalsa::{IngredientIndex, Zalsa};
1416
use crate::{Database, Id, Revision};
1517

@@ -110,6 +112,10 @@ impl<A: Accumulator> Ingredient for IngredientImpl<A> {
110112
fn debug_name(&self) -> &'static str {
111113
A::DEBUG_NAME
112114
}
115+
116+
fn memo_table_types(&self) -> Arc<MemoTableTypes> {
117+
unreachable!("accumulator does not allocate pages")
118+
}
113119
}
114120

115121
impl<A> std::fmt::Debug for IngredientImpl<A>

‎src/function.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use std::any::Any;
22
use std::fmt;
33
use std::ptr::NonNull;
4+
use std::sync::Arc;
45

56
pub(crate) use maybe_changed_after::VerifyResult;
67

@@ -11,6 +12,7 @@ use crate::ingredient::{fmt_index, Ingredient};
1112
use crate::key::DatabaseKeyIndex;
1213
use crate::plumbing::MemoIngredientMap;
1314
use crate::salsa_struct::SalsaStructInDb;
15+
use crate::table::memo::MemoTableTypes;
1416
use crate::table::sync::ClaimResult;
1517
use crate::table::Table;
1618
use crate::views::DatabaseDownCaster;
@@ -30,6 +32,8 @@ mod maybe_changed_after;
3032
mod memo;
3133
mod specify;
3234

35+
pub type Memo<C> = memo::Memo<<C as Configuration>::Output<'static>>;
36+
3337
pub trait Configuration: Any {
3438
const DEBUG_NAME: &'static str;
3539

@@ -142,7 +146,10 @@ impl<C> IngredientImpl<C>
142146
where
143147
C: Configuration,
144148
{
145-
pub fn new(
149+
/// # Safety
150+
///
151+
/// `memo_type` and `memo_table_types` must be correct.
152+
pub unsafe fn new(
146153
index: IngredientIndex,
147154
memo_ingredient_indices: <C::SalsaStruct<'static> as SalsaStructInDb>::MemoIngredientMap,
148155
lru: usize,
@@ -322,6 +329,10 @@ where
322329
C::DEBUG_NAME
323330
}
324331

332+
fn memo_table_types(&self) -> Arc<MemoTableTypes> {
333+
unreachable!("function does not allocate pages")
334+
}
335+
325336
fn cycle_recovery_strategy(&self) -> CycleRecoveryStrategy {
326337
C::CYCLE_STRATEGY
327338
}

‎src/function/memo.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use crate::cycle::{CycleHeadKind, CycleHeads, CycleRecoveryStrategy, EMPTY_CYCLE
99
use crate::function::{Configuration, IngredientImpl};
1010
use crate::key::DatabaseKeyIndex;
1111
use crate::revision::AtomicRevision;
12-
use crate::table::memo::MemoTable;
12+
use crate::table::memo::MemoTableWithTypesMut;
1313
use crate::zalsa::{MemoIngredientIndex, Zalsa};
1414
use crate::zalsa_local::{QueryOrigin, QueryRevisions};
1515
use crate::{Event, EventKind, Id, Revision};
@@ -84,7 +84,7 @@ impl<C: Configuration> IngredientImpl<C> {
8484
/// with an equivalent memo that has no value. If the memo is untracked, FixpointInitial,
8585
/// or has values assigned as output of another query, this has no effect.
8686
pub(super) fn evict_value_from_memo_for(
87-
table: &mut MemoTable,
87+
table: MemoTableWithTypesMut<'_>,
8888
memo_ingredient_index: MemoIngredientIndex,
8989
) {
9090
let map = |memo: &mut Memo<C::Output<'static>>| {
@@ -122,7 +122,7 @@ impl<C: Configuration> IngredientImpl<C> {
122122
}
123123

124124
#[derive(Debug)]
125-
pub(super) struct Memo<V> {
125+
pub struct Memo<V> {
126126
/// The result of the query, if we decide to memoize it.
127127
pub(super) value: Option<V>,
128128

‎src/ingredient.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
use std::any::{Any, TypeId};
22
use std::fmt;
3+
use std::sync::Arc;
34

45
use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues};
56
use crate::cycle::{CycleHeadKind, CycleRecoveryStrategy};
67
use crate::function::VerifyResult;
78
use crate::plumbing::IngredientIndices;
9+
use crate::table::memo::MemoTableTypes;
810
use crate::table::Table;
911
use crate::zalsa::{transmute_data_mut_ptr, transmute_data_ptr, IngredientIndex, Zalsa};
1012
use crate::zalsa_local::QueryOrigin;
@@ -132,6 +134,8 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
132134
);
133135
}
134136

137+
fn memo_table_types(&self) -> Arc<MemoTableTypes>;
138+
135139
fn fmt_index(&self, index: crate::Id, fmt: &mut fmt::Formatter<'_>) -> fmt::Result;
136140
// Function ingredient methods
137141

‎src/input.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use std::any::{Any, TypeId};
22
use std::fmt;
33
use std::ops::DerefMut;
4+
use std::sync::Arc;
45

56
pub mod input_field;
67
pub mod setter;
@@ -14,7 +15,7 @@ use crate::ingredient::{fmt_index, Ingredient};
1415
use crate::input::singleton::{Singleton, SingletonChoice};
1516
use crate::key::DatabaseKeyIndex;
1617
use crate::plumbing::{Jar, Stamp};
17-
use crate::table::memo::MemoTable;
18+
use crate::table::memo::{MemoTable, MemoTableTypes};
1819
use crate::table::sync::SyncTable;
1920
use crate::table::{Slot, Table};
2021
use crate::zalsa::{IngredientIndex, Zalsa};
@@ -72,6 +73,7 @@ impl<C: Configuration> Jar for JarImpl<C> {
7273
pub struct IngredientImpl<C: Configuration> {
7374
ingredient_index: IngredientIndex,
7475
singleton: C::Singleton,
76+
memo_table_types: Arc<MemoTableTypes>,
7577
_phantom: std::marker::PhantomData<C::Struct>,
7678
}
7779

@@ -80,6 +82,7 @@ impl<C: Configuration> IngredientImpl<C> {
8082
Self {
8183
ingredient_index: index,
8284
singleton: Default::default(),
85+
memo_table_types: Arc::new(MemoTableTypes::default()),
8386
_phantom: std::marker::PhantomData,
8487
}
8588
}
@@ -100,7 +103,7 @@ impl<C: Configuration> IngredientImpl<C> {
100103
let (zalsa, zalsa_local) = db.zalsas();
101104

102105
let id = self.singleton.with_scope(|| {
103-
zalsa_local.allocate(zalsa.table(), self.ingredient_index, |_| Value::<C> {
106+
zalsa_local.allocate(zalsa, self.ingredient_index, |_| Value::<C> {
104107
fields,
105108
stamps,
106109
memos: Default::default(),
@@ -219,6 +222,10 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
219222
fn debug_name(&self) -> &'static str {
220223
C::DEBUG_NAME
221224
}
225+
226+
fn memo_table_types(&self) -> Arc<MemoTableTypes> {
227+
self.memo_table_types.clone()
228+
}
222229
}
223230

224231
impl<C: Configuration> std::fmt::Debug for IngredientImpl<C> {

‎src/input/input_field.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
use std::fmt;
22
use std::marker::PhantomData;
3+
use std::sync::Arc;
34

45
use crate::function::VerifyResult;
56
use crate::ingredient::{fmt_index, Ingredient};
67
use crate::input::{Configuration, IngredientImpl, Value};
8+
use crate::table::memo::MemoTableTypes;
79
use crate::zalsa::IngredientIndex;
810
use crate::{Database, Id, Revision};
911

@@ -65,6 +67,10 @@ where
6567
fn debug_name(&self) -> &'static str {
6668
C::FIELD_DEBUG_NAMES[self.field_index]
6769
}
70+
71+
fn memo_table_types(&self) -> Arc<MemoTableTypes> {
72+
unreachable!("input fields do not allocate pages")
73+
}
6874
}
6975

7076
impl<C> std::fmt::Debug for FieldIngredientImpl<C>

‎src/interned.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ use crate::hash::FxDashMap;
1616
use crate::ingredient::{fmt_index, Ingredient};
1717
use crate::plumbing::{IngredientIndices, Jar};
1818
use crate::revision::AtomicRevision;
19-
use crate::table::memo::MemoTable;
19+
use crate::table::memo::{MemoTable, MemoTableTypes};
2020
use crate::table::sync::SyncTable;
2121
use crate::table::Slot;
2222
use crate::zalsa::{IngredientIndex, Zalsa};
@@ -62,6 +62,8 @@ pub struct IngredientImpl<C: Configuration> {
6262
///
6363
/// Deadlock requirement: We access `value_map` while holding lock on `key_map`, but not vice versa.
6464
key_map: FxDashMap<C::Fields<'static>, Id>,
65+
66+
memo_table_types: Arc<MemoTableTypes>,
6567
}
6668

6769
/// Struct storing the interned fields.
@@ -132,6 +134,7 @@ where
132134
Self {
133135
ingredient_index,
134136
key_map: Default::default(),
137+
memo_table_types: Arc::new(MemoTableTypes::default()),
135138
}
136139
}
137140

@@ -279,8 +282,6 @@ where
279282

280283
// We won any races so should intern the data
281284
Err(slot) => {
282-
let table = zalsa.table();
283-
284285
// Record the durability of the current query on the interned value.
285286
let (durability, last_interned_at) = zalsa_local
286287
.active_query()
@@ -289,7 +290,7 @@ where
289290
// `last_interned_at` needs to be `Revision::MAX`, see the intern_access_in_different_revision test.
290291
.unwrap_or((Durability::MAX, Revision::max()));
291292

292-
let id = zalsa_local.allocate(table, self.ingredient_index, |id| Value::<C> {
293+
let id = zalsa_local.allocate(zalsa, self.ingredient_index, |id| Value::<C> {
293294
fields: unsafe { self.to_internal_data(assemble(id, key)) },
294295
memos: Default::default(),
295296
syncs: Default::default(),
@@ -299,7 +300,7 @@ where
299300
last_interned_at: AtomicRevision::from(last_interned_at),
300301
});
301302

302-
let value = table.get::<Value<C>>(id);
303+
let value = zalsa.table().get::<Value<C>>(id);
303304

304305
let slot_value = (value.fields.clone(), SharedValue::new(id));
305306
unsafe { lock.insert_in_slot(data_hash, slot, slot_value) };
@@ -308,7 +309,7 @@ where
308309
data_hash,
309310
self.key_map
310311
.hasher()
311-
.hash_one(table.get::<Value<C>>(id).fields.clone())
312+
.hash_one(zalsa.table().get::<Value<C>>(id).fields.clone())
312313
);
313314

314315
// Record a dependency on this value.
@@ -415,6 +416,10 @@ where
415416
fn debug_name(&self) -> &'static str {
416417
C::DEBUG_NAME
417418
}
419+
420+
fn memo_table_types(&self) -> Arc<MemoTableTypes> {
421+
self.memo_table_types.clone()
422+
}
418423
}
419424

420425
impl<C> std::fmt::Debug for IngredientImpl<C>

‎src/lib.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ pub mod plumbing {
8787
pub use crate::key::DatabaseKeyIndex;
8888
pub use crate::memo_ingredient_indices::{
8989
IngredientIndices, MemoIngredientIndices, MemoIngredientMap, MemoIngredientSingletonIndex,
90+
NewMemoIngredientIndices,
9091
};
9192
pub use crate::revision::Revision;
9293
pub use crate::runtime::{stamp, Runtime, Stamp, StampedValue};
@@ -118,7 +119,10 @@ pub mod plumbing {
118119
}
119120

120121
pub mod function {
121-
pub use crate::function::{Configuration, IngredientImpl};
122+
pub use crate::function::Configuration;
123+
pub use crate::function::IngredientImpl;
124+
pub use crate::function::Memo;
125+
pub use crate::table::memo::MemoEntryType;
122126
}
123127

124128
pub mod tracked_struct {

0 commit comments

Comments
 (0)