[RFC] fixpoint iteration support #603

carljm · 2024-10-23T21:23:12Z

This PR removes the existing unwind-based cycle fallback support (a plus for WASM compatibility), and replaces it with support for fixpoint iteration of cycles.

To opt in to fixpoint iteration, provide two additional arguments to salsa::tracked on the definition of a tracked function: cycle_initial and cycle_fn. The former is a function which should provide a provisional starting value for fixpoint iteration on this query, and the latter is a function which has the opportunity, after each iteration that failed to converge, to decide whether to continue iterating or fallback to some fixed value. See the added test in cycle_fixpoint.rs for details.

Usability points that should be covered in the documentation:

With the old cycle fallback, it was sufficient to avoid panic for at least one query in a cycle to define a cycle fallback. With fixpoint iteration, to avoid cycle panics you must define cycle_fn and cycle_initial on every query that might end up as the "head" of a cycle (that is, queried for its value while it is already executing.)
It is entirely possible to define cycle_fn and cycle_initial so as to cause iteration to diverge and never terminate; it's up to the user to avoid this. Techniques to avoid this include a) ensuring that cycles will converge, by defining the initial value and the queries themselves monotonically (for example, in a type-inference scenario, the initial value is the bottom, or empty, type, and types will only widen, never narrow, as the cycle iterates -- thus the cycle must eventually converge to the top type, if nowhere else), and/or b) with a larger hammer, by ensuring that cycle_fn respects the iteration count it is given, and always halts iteration with a fallback value if the count reaches some "too large" value.
It's also entirely possible to define cycle_fn and cycle_initial such that memoized results can vary depending only on the order in which queries occur. Avoid this by minimizing the number of tracked functions that support fixpoint iteration and ensuring initial values and fallback values are consistent among tracked functions that may occur in a cycle together.

This is an RFC pull request to get initial reviewer feedback on the design and implementation. Remaining TODO items:

add tests for more complex cycles:
- nested (multiple head) cycles
- cycles with multiple paths back to the same cycle head
add tests for cross-thread cycles
add tests that use inputs in cycle recovery functions
test in red-knot and validate it works there
performance improvements
- lazy creation of initial-value memo?
documentation

netlify · 2024-10-23T21:23:37Z

✅ Deploy Preview for salsa-rs canceled.

Name	Link
🔨 Latest commit	`5202579`
🔍 Latest deploy log	https://app.netlify.com/sites/salsa-rs/deploys/676236fe2fc7440008e84cd4

codspeed-hq · 2024-10-23T21:25:09Z

CodSpeed Performance Report

Merging #603 will degrade performances by 25.11%

_{Comparing carljm:fixpoint (5202579) with master (3c7f169)}

Summary

❌ 5 regressions
✅ 4 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	`master`	`carljm:fixpoint`	Change
❌	`new[Input]`	13.8 µs	17.9 µs	-23.01%
❌	`mutating[10]`	20 µs	26.2 µs	-23.69%
❌	`mutating[20]`	21.3 µs	27.7 µs	-23.01%
❌	`mutating[30]`	21.5 µs	28.7 µs	-25.11%
❌	`many_tracked_structs`	49.2 µs	55 µs	-10.59%

nikomatsakis · 2024-10-24T01:50:42Z

This is very cool! (Admittedly, I say this pre-review.)

MichaReiser

This looks great. I left a few comments where I struggled understanding the implementation or had smaller suggestions.

MichaReiser · 2024-10-24T07:44:23Z

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

@@ -179,12 +182,16 @@ macro_rules! setup_tracked_fn {
                    $inner($db, $($input_id),*)
                }

+                fn cycle_initial<$db_lt>(db: &$db_lt dyn $Db) -> Self::Output<$db_lt> {
+                    $($cycle_recovery_initial)*(db)


I suspect that it's possible that the initial value function or the recover functions itself could create new cycles. Is this indeed the case and if so, what's salsa's behavior?

Hmm. In the use cases I have in mind (e.g. for red-knot) there should not be any reason to call a query from within one of these functions. And it seems like this could be quite difficult to deal with. So unless we have clear use cases, I would rather consider this out of scope. We could just document "don't do that," or we could add some kind of explicit prevention.

Yeah. I'm not suggesting that users should do that. I only wonder about what happens if a user does it regardless.

Ideally, we wouldn't provide a Db but we can't do that because a user might want to create a tracked struct.

I think the simplest approach here would be to add some state on Runtime that causes an error if you try to push a new active query. I'll wait for Niko review before doing that, though.

Added a TODO comment on this to get Niko's feedback.

If this replaces the old cycle handling, we (rust-analyzer) actually needs to be able to call queries here as we already do that in some of our cycle recovery functions today. I think most of those are in part because of assumptions else where in the code, an example being generic_defaults always returning a list with length matching the number of generics, and to compute the number we need to call a query. That query can run into a cycle though, so we still need to figure out the number of generics in the recovery function which means we need to call the corresponding query.

src/function/execute.rs

src/zalsa_local.rs

MichaReiser · 2024-10-24T08:03:34Z

tests/cycle_fixpoint.rs

+    }
+}
+
+#[salsa::tracked(cycle_fn=cycle_recover, cycle_initial=cycle_initial)]


Should we define a custom CycleRecovery trait that defines recover and initial methods. It would give us a good point to put cycle handling documentation and avoids the risk for incorrect-macro use when only specifiying one but not boht values.

Yes, I thought about this as well. It's not clear to me what is actually the better UX. Implementing a trait is somewhat more boilerplate, and in practice there's not much difference in the experience if you fail to implement one of the methods. Either way the compiler catches the error for you, because its an error in macro expansion if you don't provide both cycle_fn and cycle_initial, with a trait it would be a compiler error that you didn't fully implement the trait.

I guess with a trait your IDE might give you the right signatures of the methods for free, which is kind of nice...

I normally don't like traits, but I feel like a trait would be preferable because the current annotation seems a little too wordy for my tastes.

Ok, going to still wait for Niko's feedback on this point before updating, but the trait idea makes sense to me.

Actually, I'm okay with a non-trait based approach. rust-analyzer uses a #[salsa::cycle(path::to::function)]-style annotation, so I think to ease the transition, avoiding a trait would be nice.

rust-analyzer uses a #[salsa::cycle(path::to::function)]-style annotation, so I think to ease the transition, avoiding a trait would be nice.

Is cycle handling common in ra? I do see how it reduces the diff size because it isn't necessary to make the function a trait-method but you would still have to change every use because you now have to specify both the cycle and cycle initial functions.

Added a TODO comment in the code for this question as well.

its common but not common enough that it would matter (there are ~11 cycle handlers and the majority of them return a dummy result)

One theoreitcal nice thing about taking the functions here directly would be that we could allow supporting closure expressions in that position inline instead of having to create a whole function (as cycles that are not going to be fixpoint are likely not gonna have complex functions here), a trait would be quite wordy imo.

carljm · 2024-10-29T18:15:08Z

In writing more comprehensive tests for this, I realized that it needs some changes to correctly handle multi-revision scenarios; taking it to Draft mode until I get that fixed.

carljm · 2024-10-30T00:37:43Z

Ok, multiple-revision cases are now fixed, and we now populate the initial provisional value only lazily, in case a cycle is actually encountered, which should reduce the number of memos created by quite a lot.

Also added a bunch of tests, including multiple-revision cases and one test involving durability. Still need to add cross-thread cycle tests.

carljm · 2024-10-30T00:42:47Z

tests/cycle/main.rs

+// Diagram nomenclature for nodes: Each node is represented as a:xx(ii), where `a` is a sequential
+// identifier from `a`, `b`, `c`..., xx is one of the four query kinds:
+// - `Ni` for `min_iterate`
+// - `Xi` for `max_iterate`
+// - `Np` for `min_panic`
+// - `Xp` for `max_panic`
+//
+// and `ii` is the inputs for that query, represented as a comma-separated list, with each
+// component representing an input:
+// - `a`, `b`, `c`... where the input is another node,
+// - `uXX` for `UntrackedRead(XX)`
+// - `vXX` for `Value(XX)`
+// - `sY` for `Successor(Y)`
+//


These are admittedly obscure-looking chicken scratches. I'm not claiming they are super readable, but they are concise enough to put into an ASCII graph diagram, and (once you get familiar with them) give a lot of information about the behavior of the test. They were really helpful to me in writing and debugging the tests.

Open to feedback that I should do this differently for better readability by future maintainers...

src/function/execute.rs

MichaReiser

The lazy creation of the initial value is a neat improvement. Nice for taking the time to work on it !

src/function/execute.rs

MichaReiser · 2024-10-30T07:35:48Z

The benchmarks show a 4-5% regression. It seems that we're now resizing some hash maps more often. Are we reporting more tracked reads than before? Could you take a look what's causing it?

carljm · 2024-11-01T00:42:07Z

Initial experiments using this in the red-knot type checker are promising: astral-sh/ruff#14029

Not yet using it for loopy control flow in that PR, but there are cycles in the core type definitions of Python builtins and standard library, which we previously had a hacky fallback in place for using Salsa's previous cycle fallback support. Moving over to fixpoint iteration just worked, and fixed the type of a builtin impacted by the cycle.

On the downside, it is a performance regression. Need to do more work there.

MichaReiser · 2024-11-15T07:39:34Z

src/function/maybe_changed_after.rs

@@ -73,22 +114,29 @@ where
        );

        // Check if the inputs are still valid and we can just compare `changed_at`.
-        if self.deep_verify_memo(db, &old_memo, &active_query) {
-            return Some(old_memo.revisions.changed_at > revision);
+        let active_query = zalsa_local.push_query(database_key_index);


From looking at the red knot benchmarks, the regression mainly comes from the extra push_query calls here (that probably also applies for queries not participating in cycles) and the constructed hash set in deep_verify_memo (specific to red knot?)

Could we move the push_query call into deep_verify_memo after the second shallow_verify_memo or does that result in deadlocks? Just so that we can avoid pushing queries unless it's absolutely necessary

That seems weird, this isnt an extra push_query call after all. This merely moved down a couple lines

Edit: Nevermind, there are more in execute

MichaReiser · 2024-11-15T07:42:09Z

src/function/maybe_changed_after.rs

-                            {
-                                return false;
+        loop {
+            let mut cycle_heads = FxHashSet::default();


Is it intentional that we create a new cycle_heads in every iteration? Could we re-use the cycle_heads and instead call clear to avoid re-allocating the hash set on every iteration?

MichaReiser · 2024-11-17T17:14:04Z

src/ingredient.rs

@@ -38,7 +39,7 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
        db: &'db dyn Database,
        input: Option<Id>,
        revision: Revision,
-    ) -> bool;
+    ) -> VerifyResult;


I think this change will also come handy for the "faster accumulator" work because we'll need to also return whether the ingredient had any accumulated values.

Veykril · 2024-12-28T15:50:36Z

src/cycle.rs

-        })
-    }
+    /// Cut off iteration and use the given result value for this query.
+    Fallback(T),


Just for my own understanding, but having a cycle_fn that returns Fallback immediately would effectively be the same as the previous cycle handling right? And it would never call the initial function then.

In the current implementation, the initial function will still be called. If query A(0) is executing, it calls B(0), which in turn calls A(0), the second call to A(0) will call the initial function to get an initial value to return to B(0), which will compute its provisional value and return that to A(0), which will then call the cycle recovery function. If that cycle recovery function immediately returns Fallback, then that fallback value will be used for A(0). (I think we also need to then iterate the cycle one more time in order to populate a correct value for B(0) based on the fallback value for A(0), but that's not currently implemented.)

It seems like it could be possible to instead call the cycle recovery function first thing when we encounter the cycle, and allow it to return Fallback immediately, before the initial function is ever called. That seems worth experimenting with.

A key difference in this PR from the current cycle handling is that in the current cycle handling all cycle participants that have a cycle fallback receive a fallback value; this ensures determinism regardless of query order, but it means that you can't have some participants in the cycle compute a real value based on a fallback for some other cycle participant.

In this PR, if your recovery function falls back to some value, that fallback only applies to the individual query, and the rest of the cycle will then resolve normally using that fallback value.

I see, thanks, that explanation made this very clear!

Veykril · 2024-12-28T15:55:38Z

src/table/sync.rs

+pub(crate) enum ClaimResult<'a> {
+    Retry,
+    Cycle,
+    Claimed(ClaimGuard<'a>),
+}


nit: ClamGuard only has one niche so this might affect perf slightly compared to the Option wrapping from before

Funnily enough, we can give it 254 niches by adding a bool field to it which occupies its padding, something to keep in mind.

Veykril · 2024-12-28T16:16:43Z

src/function/execute.rs

+        // only when a cycle is actually encountered.
+        let mut opt_last_provisional: Option<&Memo<<C as Configuration>::Output<'db>>> = None;
+
+        loop {


Can we branch on the cycle strategy here? There is a lot of work that I don't think needs to happen for Panic strategies that non cycle supporting queries currently pay (and given the branching would happen on a const it should optimized out)

Veykril · 2024-12-28T16:23:06Z

src/function/maybe_changed_after.rs

+/// Result of memo validation.
+pub enum VerifyResult {
+    /// Memo has changed and needs to be recomputed.
+    Changed,
+
+    /// Memo remains valid.
+    ///
+    /// Database keys in the hashset represent cycle heads encountered in validation; don't mark
+    /// memos verified until we've iterated the full cycle to ensure no inputs changed.
+    Unchanged(FxHashSet<DatabaseKeyIndex>),
+}


This change (bool being replaced by VerifyResult) also seems like it will make non cycle queries pay a hefty perf price. That type has a destructor and is 40 bytes where as the previous return value could fit into a register, this can't.

I'll focus first on getting consistently right behavior (including across threads) and then revisit this and other possible ways to claw back perf. Off the top of my head I'm not sure how to improve this, will need to think about it.

Yes of course, I am mainly noting down what I saw, no need to tackle perf nits/reviews until things are working properly. (also happy to help out with perf work on this PR once its at that stage)

Veykril · 2024-12-29T10:00:00Z

On a related perf note, would it make sense to move the dataflow test (or a maybe slightly more involved version) into a benchmark?

carljm requested review from nikomatsakis, MichaReiser and davidbarsky October 23, 2024 21:23

MichaReiser reviewed Oct 24, 2024

View reviewed changes

carljm force-pushed the fixpoint branch from f700c39 to f411766 Compare October 24, 2024 20:06

salsa-rs deleted a comment from MichaReiser Oct 24, 2024

carljm marked this pull request as draft October 29, 2024 18:15

carljm marked this pull request as ready for review October 30, 2024 00:37

carljm commented Oct 30, 2024

View reviewed changes

MichaReiser reviewed Oct 30, 2024

View reviewed changes

src/function/execute.rs Outdated Show resolved Hide resolved

MichaReiser reviewed Oct 30, 2024

View reviewed changes

src/function/execute.rs Outdated Show resolved Hide resolved

src/function/execute.rs Outdated Show resolved Hide resolved

src/function/execute.rs Outdated Show resolved Hide resolved

carljm force-pushed the fixpoint branch from c1bbdcf to f44f2f7 Compare November 14, 2024 23:53

MichaReiser reviewed Nov 15, 2024

View reviewed changes

MichaReiser reviewed Nov 17, 2024

View reviewed changes

carljm added 8 commits December 16, 2024 15:35

add example test for fixpoint iteration

3d00ae3

add a multi-symbol test

31979b9

simplify test case

35e236e

WIP: remove existing cycle handling tests for now

90ea6e9

WIP: remove all existing cycle handling, add fixpoint options

f218e58

WIP: added provisional value and cycle fields

6ce61c6

rename to CycleRecoveryStrategy::Fixpoint

29d110e

WIP: rip out ProvisionalValue

adb6c77

carljm and others added 26 commits December 16, 2024 15:35

WIP: working single-iteration with provisional memo

7f82b6d

WIP: add count arg to cycle_recovery_fn

a7be3d9

WIP: move insert-initial out into fetch_cold

6c6dd55

WIP: cycle-head iteration

8974361

WIP: move loop into execute

797655e

WIP: delay storing memo

27742a2

WIP: remove ourself from cycle heads when done iterating

4f4df72

WIP: working convergence and fallback

315944e

WIP: clippy and cleanup

28242a4

WIP: improve comments and add a type annotation

0be825a

WIP: don't allow cycle_fn with no_eq

c3c84c4

WIP: add tracing for cycle iteration

72dff5f

WIP: fail fast if we get an evicted provisional value

6b44c92

WIP: use FxHashSet::from_iter

a029ef4

add tests, fix multiple-revision, lazy provisional value

492ae1b

review feedback, more tracing

0f7d940

fix multi-revision bug

5ef7a5f

better fix for multi-revision bug

3093460

test fixes

7d9ec1c

pass inputs to cycle recovery functions

aa4a731

fixed cycle-unchanged test

67376f1

add TODO comments for some outstanding questions

286b5fb

add a test for the "AB peeping C" scenario

b2d4d92

another parallel test scenario

670f88b

WIP: removed cycle_ignore; nested cycles broken

00acc56

fixed all single-thread cycles; multi-thread still not working

5202579

carljm force-pushed the fixpoint branch from 8325b62 to 5202579 Compare December 18, 2024 02:44

Veykril reviewed Dec 28, 2024

View reviewed changes

[RFC] fixpoint iteration support #603

Are you sure you want to change the base?

[RFC] fixpoint iteration support #603

Conversation

carljm commented Oct 23, 2024 • edited Loading

netlify bot commented Oct 23, 2024 • edited Loading

✅ Deploy Preview for salsa-rs canceled.

codspeed-hq bot commented Oct 23, 2024 • edited Loading

Merging #603 will degrade performances by 25.11%

Summary

Benchmarks breakdown

nikomatsakis commented Oct 24, 2024

MichaReiser left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davidbarsky Oct 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

carljm commented Oct 29, 2024

carljm commented Oct 30, 2024

Choose a reason for hiding this comment

MichaReiser left a comment

Choose a reason for hiding this comment

MichaReiser commented Oct 30, 2024 • edited Loading

carljm commented Nov 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Veykril Jan 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Veykril Dec 29, 2024 • edited Loading

Choose a reason for hiding this comment

Veykril commented Dec 29, 2024

carljm commented Oct 23, 2024 •

edited

Loading

netlify bot commented Oct 23, 2024 •

edited

Loading

codspeed-hq bot commented Oct 23, 2024 •

edited

Loading

davidbarsky Oct 24, 2024 •

edited

Loading

MichaReiser commented Oct 30, 2024 •

edited

Loading

Veykril Jan 5, 2025 •

edited

Loading

Veykril Dec 29, 2024 •

edited

Loading