rewrite cycle handling to support fixed-point iteration (#603)

Author: carljm

Cargo.toml

@@ -73,6 +73,10 @@ harness = false
 name = "accumulator"
 harness = false
 
+[[bench]]
+name = "dataflow"
+harness = false
+
 [workspace]
 members = ["components/salsa-macro-rules", "components/salsa-macros"]
 

benches/dataflow.rs

@@ -0,0 +1,170 @@
+//! Benchmark for fixpoint iteration cycle resolution.
+//!
+//! This benchmark simulates a (very simplified) version of a real dataflow analysis using fixpoint
+//! iteration.
+use codspeed_criterion_compat::{criterion_group, criterion_main, BatchSize, Criterion};
+use salsa::{CycleRecoveryAction, Database as Db, Setter};
+use std::collections::BTreeSet;
+use std::iter::IntoIterator;
+
+/// A Use of a symbol.
+#[salsa::input]
+struct Use {
+    reaching_definitions: Vec<Definition>,
+}
+
+/// A Definition of a symbol, either of the form `base + increment` or `0 + increment`.
+#[salsa::input]
+struct Definition {
+    base: Option<Use>,
+    increment: usize,
+}
+
+#[derive(Eq, PartialEq, Clone, Debug, salsa::Update)]
+enum Type {
+    Bottom,
+    Values(Box<[usize]>),
+    Top,
+}
+
+impl Type {
+    fn join(tys: impl IntoIterator<Item = Type>) -> Type {
+        let mut result = Type::Bottom;
+        for ty in tys.into_iter() {
+            result = match (result, ty) {
+                (result, Type::Bottom) => result,
+                (_, Type::Top) => Type::Top,
+                (Type::Top, _) => Type::Top,
+                (Type::Bottom, ty) => ty,
+                (Type::Values(a_ints), Type::Values(b_ints)) => {
+                    let mut set = BTreeSet::new();
+                    set.extend(a_ints);
+                    set.extend(b_ints);
+                    Type::Values(set.into_iter().collect())
+                }
+            }
+        }
+        result
+    }
+}
+
+#[salsa::tracked(cycle_fn=use_cycle_recover, cycle_initial=use_cycle_initial)]
+fn infer_use<'db>(db: &'db dyn Db, u: Use) -> Type {
+    let defs = u.reaching_definitions(db);
+    match defs[..] {
+        [] => Type::Bottom,
+        [def] => infer_definition(db, def),
+        _ => Type::join(defs.iter().map(|&def| infer_definition(db, def))),
+    }
+}
+
+#[salsa::tracked(cycle_fn=def_cycle_recover, cycle_initial=def_cycle_initial)]
+fn infer_definition<'db>(db: &'db dyn Db, def: Definition) -> Type {
+    let increment_ty = Type::Values(Box::from([def.increment(db)]));
+    if let Some(base) = def.base(db) {
+        let base_ty = infer_use(db, base);
+        add(&base_ty, &increment_ty)
+    } else {
+        increment_ty
+    }
+}
+
+fn def_cycle_initial(_db: &dyn Db, _def: Definition) -> Type {
+    Type::Bottom
+}
+
+fn def_cycle_recover(
+    _db: &dyn Db,
+    value: &Type,
+    count: u32,
+    _def: Definition,
+) -> CycleRecoveryAction<Type> {
+    cycle_recover(value, count)
+}
+
+fn use_cycle_initial(_db: &dyn Db, _use: Use) -> Type {
+    Type::Bottom
+}
+
+fn use_cycle_recover(
+    _db: &dyn Db,
+    value: &Type,
+    count: u32,
+    _use: Use,
+) -> CycleRecoveryAction<Type> {
+    cycle_recover(value, count)
+}
+
+fn cycle_recover(value: &Type, count: u32) -> CycleRecoveryAction<Type> {
+    match value {
+        Type::Bottom => CycleRecoveryAction::Iterate,
+        Type::Values(_) => {
+            if count > 4 {
+                CycleRecoveryAction::Fallback(Type::Top)
+            } else {
+                CycleRecoveryAction::Iterate
+            }
+        }
+        Type::Top => CycleRecoveryAction::Iterate,
+    }
+}
+
+fn add(a: &Type, b: &Type) -> Type {
+    match (a, b) {
+        (Type::Bottom, _) | (_, Type::Bottom) => Type::Bottom,
+        (Type::Top, _) | (_, Type::Top) => Type::Top,
+        (Type::Values(a_ints), Type::Values(b_ints)) => {
+            let mut set = BTreeSet::new();
+            set.extend(
+                a_ints
+                    .into_iter()
+                    .flat_map(|a| b_ints.into_iter().map(move |b| a + b)),
+            );
+            Type::Values(set.into_iter().collect())
+        }
+    }
+}
+
+fn dataflow(criterion: &mut Criterion) {
+    criterion.bench_function("converge_diverge", |b| {
+        b.iter_batched_ref(
+            || {
+                let mut db = salsa::DatabaseImpl::new();
+
+                let defx0 = Definition::new(&db, None, 0);
+                let defy0 = Definition::new(&db, None, 0);
+                let defx1 = Definition::new(&db, None, 0);
+                let defy1 = Definition::new(&db, None, 0);
+                let use_x = Use::new(&db, vec![defx0, defx1]);
+                let use_y = Use::new(&db, vec![defy0, defy1]);
+                defx1.set_base(&mut db).to(Some(use_y));
+                defy1.set_base(&mut db).to(Some(use_x));
+
+                // prewarm cache
+                let _ = infer_use(&db, use_x);
+                let _ = infer_use(&db, use_y);
+
+                (db, defx1, use_x, use_y)
+            },
+            |(db, defx1, use_x, use_y)| {
+                // Set the increment on x to 0.
+                defx1.set_increment(db).to(0);
+
+                // Both symbols converge on 0.
+                assert_eq!(infer_use(db, *use_x), Type::Values(Box::from([0])));
+                assert_eq!(infer_use(db, *use_y), Type::Values(Box::from([0])));
+
+                // Set the increment on x to 1.
+                defx1.set_increment(db).to(1);
+
+                // Now the loop diverges and we fall back to Top.
+                assert_eq!(infer_use(db, *use_x), Type::Top);
+                assert_eq!(infer_use(db, *use_y), Type::Top);
+            },
+            BatchSize::LargeInput,
+        );
+    });
+}
+
+criterion_group!(benches, dataflow);
+criterion_main!(benches);

book/src/SUMMARY.md

@@ -22,7 +22,6 @@
   - [On-demand (Lazy) inputs](./common_patterns/on_demand_inputs.md)
 - [Tuning](./tuning.md)
 - [Cycle handling](./cycles.md)
-  - [Recovering via fallback](./cycles/fallback.md)
 
 # How Salsa works internally
 

book/src/cycles.md

@@ -1,5 +1,40 @@
 # Cycle handling
 
-By default, when Salsa detects a cycle in the computation graph, Salsa will panic with a [`salsa::Cycle`] as the panic value. The [`salsa::Cycle`] structure that describes the cycle, which can be useful for diagnosing what went wrong.
+By default, when Salsa detects a cycle in the computation graph, Salsa will panic with a message naming the "cycle head"; this is the query that was called while it was also on the active query stack, creating a cycle.
 
-[`salsa::cycle`]: https://github.com/salsa-rs/salsa/blob/0f9971ad94d5d137f1192fde2b02ccf1d2aca28c/src/lib.rs#L654-L672
+Salsa also supports recovering from query cycles via fixed-point iteration. Fixed-point iteration is only usable if the queries which may be involved in a cycle are monotone and operate on a value domain which is a partial order with fixed height. Effectively, this means that the queries' output must always be "larger" than its input, and there must be some "maximum" or "top" value. This ensures that fixed-point iteration will converge to a value. (A typical case would be queries operating on types, which form a partial order with a "top" type.)
+
+In order to support fixed-point iteration for a query, provide the `cycle_fn` and `cycle_initial` arguments to `salsa::tracked`:
+
+```rust
+#[salsa::tracked(cycle_fn=cycle_fn, cycle_initial=initial_fn)]
+fn query(db: &dyn salsa::Database) -> u32 {
+    // ...
+}
+
+fn cycle_fn(_db: &dyn KnobsDatabase, _value: &u32, _count: u32) -> salsa::CycleRecoveryAction<u32> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+fn initial(_db: &dyn KnobsDatabase) -> u32 {
+    0
+}
+```
+
+If `query` becomes the head of a cycle (that is, `query` is executing and on the active query stack, it calls `query2`, `query2` calls `query3`, and `query3` calls `query` again -- there could be any number of queries involved in the cycle), the `initial_fn` will be called to generate an "initial" value for `query` in the fixed-point computation. (The initial value should usually be the "bottom" value in the partial order.) All queries in the cycle will compute a provisional result based on this initial value for the cycle head. That is, `query3` will compute a provisional result using the initial value for `query`, `query2` will compute a provisional result using this provisional value for `query3`. When `cycle2` returns its provisional result back to `cycle`, `cycle` will observe that it has received a provisional result from its own cycle, and will call the `cycle_fn` (with the current value and the number of iterations that have occurred so far). The `cycle_fn` can return `salsa::CycleRecoveryAction::Iterate` to indicate that the cycle should iterate again, or `salsa::CycleRecoveryAction::Fallback(value)` to indicate that the cycle should stop iterating and fall back to the value provided.
+
+If the `cycle_fn` continues to return `Iterate`, the cycle will iterate until it converges: that is, until two successive iterations produce the same result.
+
+If the `cycle_fn` returns `Fallback`, the cycle will iterate one last time and verify that the returned value is the same as the fallback value; that is, the fallback value results in a stable converged cycle. If not, Salsa will panic. It is not permitted to use a fallback value that does not converge, because this would leave the cycle in an unpredictable state, depending on the order of query execution.
+
+## All potential cycle heads must set `cycle_fn` and `cycle_initial`
+
+Consider a two-query cycle where `query_a` calls `query_b`, and `query_b` calls `query_a`. If `query_a` is called first, then it will become the "cycle head", but if `query_b` is called first, then `query_b` will be the cycle head. In order for a cycle to use fixed-point iteration instead of panicking, the cycle head must set `cycle_fn` and `cycle_initial`. This means that in order to be robust against varying query execution order, both `query_a` and `query_b` must set `cycle_fn` and `cycle_initial`.
+
+## Ensuring convergence
+
+Fixed-point iteration is a powerful tool, but is also easy to misuse, potentially resulting in infinite iteration. To avoid this, ensure that all queries participating in fixpoint iteration are deterministic and monotone.
+
+## Calling Salsa queries from within `cycle_fn` or `cycle_initial`
+
+It is permitted to call other Salsa queries from within the `cycle_fn` and `cycle_initial` functions. However, if these functions re-enter the same cycle, this can lead to unpredictable results. Take care which queries are called from within cycle-recovery functions, and avoid triggering further cycles.

book/src/cycles/fallback.md

@@ -15,51 +15,3 @@ When a thread `T1` attempts to execute a query `Q`, it will try to load the valu
 * Otherwise, if `Q` is being computed by some other thread `T2`, we have to check whether `T2` is (transitively) blocked on `T1`. If so, there is a cycle.
 
 These two cases are handled internally by the `Runtime::try_block_on` function. Detecting the intra-thread cycle case is easy; to detect cross-thread cycles, the runtime maintains a dependency DAG between threads (identified by `RuntimeId`). Before adding an edge `T1 -> T2` (i.e., `T1` is blocked waiting for `T2`) into the DAG, it checks whether a path exists from `T2` to `T1`. If so, we have a cycle and the edge cannot be added (then the DAG would not longer be acyclic).
-
-When a cycle is detected, the current thread `T1` has full access to the query stacks that are participating in the cycle. Consider: naturally, `T1` has access to its own stack. There is also a path `T2 -> ... -> Tn -> T1` of blocked threads. Each of the blocked threads `T2 ..= Tn` will have moved their query stacks into the dependency graph, so those query stacks are available for inspection.
-
-Using the available stacks, we can create a list of cycle participants `Q0 ... Qn` and store that into a `Cycle` struct. If none of the participants `Q0 ... Qn` have cycle recovery enabled, we panic with the `Cycle` struct, which will trigger all the queries on this thread to panic.
-
-## Cycle recovery via fallback
-
-If any of the cycle participants `Q0 ... Qn` has cycle recovery set, we recover from the cycle. To help explain how this works, we will use this example cycle which contains three threads. Beginning with the current query, the cycle participants are `QA3`, `QB2`, `QB3`, `QC2`, `QC3`, and `QA2`.
-
-```
-        The cyclic
-        edge we have
-        failed to add.
-          :
-   A      :    B         C
-          :
-   QA1    v    QB1       QC1
-┌► QA2    ┌──► QB2   ┌─► QC2
-│  QA3 ───┘    QB3 ──┘   QC3 ───┐
-│                               │
-└───────────────────────────────┘
-```
-
-Recovery works in phases:
-
-* **Analyze:** As we enumerate the query participants, we collect their collective inputs (all queries invoked so far by any cycle participant) and the max changed-at and min duration. We then remove the cycle participants themselves from this list of inputs, leaving only the queries external to the cycle.
-* **Mark**: For each query Q that is annotated with `#[salsa::cycle]`, we mark it and all of its successors on the same thread by setting its `cycle` flag to the `c: Cycle` we constructed earlier; we also reset its inputs to the collective inputs gathering during analysis. If those queries resume execution later, those marks will trigger them to immediately unwind and use cycle recovery, and the inputs will be used as the inputs to the recovery value.
-    * Note that we mark *all* the successors of Q on the same thread, whether or not they have recovery set. We'll discuss later how this is important in the case where the active thread (A, here) doesn't have any recovery set.
-* **Unblock**: Each blocked thread T that has a recovering query is forcibly reawoken; the outgoing edge from that thread to its successor in the cycle is removed. Its condvar is signalled with a `WaitResult::Cycle(c)`. When the thread reawakens, it will see that and start unwinding with the cycle `c`.
-* **Handle the current thread:** Finally, we have to choose how to have the current thread proceed. If the current thread includes any cycle with recovery information, then we can begin unwinding. Otherwise, the current thread simply continues as if there had been no cycle, and so the cyclic edge is added to the graph and the current thread blocks. This is possible because some other thread had recovery information and therefore has been awoken.
-
-Let's walk through the process with a few examples.
-
-### Example 1: Recovery on the detecting thread
-
-Consider the case where only the query QA2 has recovery set. It and QA3 will be marked with their `cycle` flag set to `c: Cycle`. Threads B and C will not be unblocked, as they do not have any cycle recovery nodes. The current thread (Thread A) will initiate unwinding with the cycle `c` as the value. Unwinding will pass through QA3 and be caught by QA2. QA2 will substitute the recovery value and return normally. QA1 and QC3 will then complete normally and so forth, on up until all queries have completed.
-
-### Example 2: Recovery in two queries on the detecting thread
-
-Consider the case where both query QA2 and QA3 have recovery set. It proceeds the same Example 1 until the current thread initiates unwinding, as described in Example 1. When QA3 receives the cycle, it stores its recovery value and completes normally. QA2 then adds QA3 as an input dependency: at that point, QA2 observes that it too has the cycle mark set, and so it initiates unwinding. The rest of QA2 therefore never executes. This unwinding is caught by QA2's entry point and it stores the recovery value and returns normally. QA1 and QC3 then continue normally, as they have not had their `cycle` flag set.
-
-### Example 3: Recovery on another thread
-
-Now consider the case where only the query QB2 has recovery set. It and QB3 will be marked with the cycle `c: Cycle` and thread B will be unblocked; the edge `QB3 -> QC2` will be removed from the dependency graph. Thread A will then add an edge `QA3 -> QB2` and block on thread B. At that point, thread A releases the lock on the dependency graph, and so thread B is re-awoken. It observes the `WaitResult::Cycle` and initiates unwinding. Unwinding proceeds through QB3 and into QB2, which recovers. QB1 is then able to execute normally, as is QA3, and execution proceeds from there.
-
-### Example 4: Recovery on all queries
-
-Now consider the case where all the queries have recovery set. In that case, they are all marked with the cycle, and all the cross-thread edges are removed from the graph. Each thread will independently awaken and initiate unwinding. Each query will recover.