alkahest_cas/simplify/
dispatch.rs

1//! Picks a parallel simplification strategy from the shape of the expression.
2//!
3//! Feature-gated behind `--features parallel`.
4//!
5//! [`super::parallel::simplify_par`] and [`super::redex::simplify_redex`] win on
6//! different shapes, and the difference is large — up to 4x either way — so the
7//! choice matters more than tuning either one.  Fork-join keeps each chain on a
8//! single worker and wins when the expression is wide, because a whole subtree
9//! stays in one core's cache.  Level scheduling wins when it is deep, because
10//! fork-join only forks on `Add`/`Mul` nodes with four or more children and a
11//! deep chain gives it nothing to fork on.
12//!
13//! The advantage is also conditional on having cores to spend: at four workers
14//! level scheduling was faster on every shape measured, including the widest,
15//! so fork-join is only considered from [`MIN_WORKERS_FOR_FORK_JOIN`] up.
16//!
17//! The shape discriminator is average level width — nodes divided by height:
18//!
19//! | shape | nodes / height | pick |
20//! |---|---|---|
21//! | deep chain (2000 levels, width 1) | ~1 | level-scheduled |
22//! | wide sum, independent terms (1024 x 8) | ~980 | fork-join |
23//! | many medium chains (1024 x 32) | ~1000 | fork-join |
24//!
25//! # Caveat on the threshold
26//!
27//! [`WIDTH_THRESHOLD`] is calibrated on synthetic shapes, not on a profile of
28//! real workloads.  The extremes it separates are unambiguous — a chain and a
29//! wide sum differ by three orders of magnitude in average width — but where
30//! exactly to cut between them is a guess, and expressions near the boundary
31//! are close enough in cost that the choice matters little either way.  Treat
32//! the constant as provisional until it can be checked against real traces.
33
34#![cfg(feature = "parallel")]
35
36use crate::deriv::log::DerivedExpr;
37use crate::kernel::{ExprData, ExprId, ExprPool};
38use crate::simplify::engine::SimplifyConfig;
39
40/// Average level width at or above which fork-join is preferred.
41///
42/// Fork-join only parallelises `Add`/`Mul` nodes with at least four children,
43/// so it needs real width to beat level scheduling.
44pub const WIDTH_THRESHOLD: f64 = 8.0;
45
46/// Worker count below which level scheduling is preferred regardless of shape.
47///
48/// Fork-join's advantage on wide expressions comes from keeping a whole subtree
49/// in one core's cache, and it only pays once there are enough cores to cover
50/// the width. Measured on a 32-core machine, level scheduling was faster on
51/// *every* shape at four workers — including the widest — while fork-join won
52/// the wide shapes at sixteen and above.
53pub const MIN_WORKERS_FOR_FORK_JOIN: usize = 8;
54
55/// Which parallel simplifier [`simplify_auto`] selected.
56#[derive(Clone, Copy, PartialEq, Eq, Debug)]
57pub enum Strategy {
58    /// [`super::parallel::simplify_par`] — recursive fork-join.
59    ForkJoin,
60    /// [`super::redex::simplify_redex`] — level-scheduled redex bag.
61    LevelScheduled,
62}
63
64/// Simplify with whichever parallel strategy suits the expression's shape.
65///
66/// Results are identical to [`crate::simplify::simplify`] either way; only the
67/// schedule differs.  Note that the derivation log is deterministic only when
68/// [`Strategy::LevelScheduled`] is chosen.
69pub fn simplify_auto(expr: ExprId, pool: &ExprPool) -> DerivedExpr<ExprId> {
70    simplify_auto_with_config(expr, pool, &SimplifyConfig::default())
71}
72
73/// Like [`simplify_auto`] but with a custom [`SimplifyConfig`].
74pub fn simplify_auto_with_config(
75    expr: ExprId,
76    pool: &ExprPool,
77    config: &SimplifyConfig,
78) -> DerivedExpr<ExprId> {
79    match choose_strategy(expr, pool) {
80        Strategy::ForkJoin => super::parallel::simplify_par_with_config(expr, pool, config),
81        Strategy::LevelScheduled => super::redex::simplify_redex_with_config(expr, pool, config),
82    }
83}
84
85/// The strategy [`simplify_auto`] would use for `expr`.
86///
87/// Exposed so callers can log or override the decision, and so tests can check
88/// it without timing anything.
89pub fn choose_strategy(expr: ExprId, pool: &ExprPool) -> Strategy {
90    // Below this many workers level scheduling wins whatever the shape, so
91    // skip the shape probe entirely rather than pay a traversal to learn
92    // something that cannot change the answer.
93    if rayon::current_num_threads() < MIN_WORKERS_FOR_FORK_JOIN {
94        return Strategy::LevelScheduled;
95    }
96    let (nodes, height) = shape(expr, pool);
97    let average_width = nodes as f64 / height.max(1) as f64;
98    if average_width >= WIDTH_THRESHOLD {
99        Strategy::ForkJoin
100    } else {
101        Strategy::LevelScheduled
102    }
103}
104
105/// Count the distinct nodes reachable from `root` and the height of the DAG.
106///
107/// Iterative, so it cannot overflow the stack on the deep expressions this is
108/// meant to detect.  Costs one traversal, which is small next to the many
109/// rule-matching passes that follow.
110fn shape(root: ExprId, pool: &ExprPool) -> (usize, u32) {
111    let n = pool.len();
112    let mut height = vec![u32::MAX; n];
113    let mut pushed = vec![false; n];
114    let mut stack: Vec<(ExprId, bool)> = vec![(root, false)];
115    let mut nodes = 0_usize;
116    let mut max_height = 0_u32;
117
118    while let Some((id, expanded)) = stack.pop() {
119        let i = id.0 as usize;
120        if expanded {
121            let h = pool.with(id, |data| {
122                let mut h = 0_u32;
123                for_each_child(data, |c| {
124                    let ch = height[c.0 as usize];
125                    debug_assert_ne!(ch, u32::MAX, "child measured after its parent");
126                    h = h.max(ch.saturating_add(1));
127                });
128                h
129            });
130            height[i] = h;
131            max_height = max_height.max(h);
132            nodes += 1;
133            continue;
134        }
135        if pushed[i] {
136            continue;
137        }
138        pushed[i] = true;
139        stack.push((id, true));
140        pool.with(id, |data| {
141            for_each_child(data, |c| {
142                if !pushed[c.0 as usize] {
143                    stack.push((c, false));
144                }
145            })
146        });
147    }
148
149    // Height counts edges; a single node is one level deep.
150    (nodes, max_height + 1)
151}
152
153/// Children a rewrite pass descends into, matching `simplify_children`.
154fn for_each_child(data: &ExprData, mut f: impl FnMut(ExprId)) {
155    match data {
156        ExprData::Add(args) | ExprData::Mul(args) => args.iter().copied().for_each(f),
157        ExprData::Func { args, .. } | ExprData::Predicate { args, .. } => {
158            args.iter().copied().for_each(f)
159        }
160        ExprData::Pow { base, exp } => {
161            f(*base);
162            f(*exp);
163        }
164        ExprData::Piecewise { branches, default } => {
165            branches.iter().for_each(|&(_, v)| f(v));
166            f(*default);
167        }
168        ExprData::Forall { body, .. } | ExprData::Exists { body, .. } => f(*body),
169        ExprData::BigO(arg) => f(*arg),
170        _ => {}
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use crate::kernel::Domain;
178    use crate::simplify::simplify;
179
180    fn p() -> ExprPool {
181        ExprPool::new()
182    }
183
184    fn junk(pool: &ExprPool, x: ExprId, depth: usize) -> ExprId {
185        let one = pool.integer(1_i32);
186        let zero = pool.integer(0_i32);
187        let mut e = x;
188        for _ in 0..depth {
189            e = pool.mul(vec![e, one]);
190            e = pool.add(vec![e, zero]);
191        }
192        e
193    }
194
195    /// Run inside a pool with enough workers that `choose_strategy` is deciding
196    /// on shape rather than falling back to the low-worker rule.
197    fn with_workers<R: Send>(f: impl FnOnce() -> R + Send) -> R {
198        rayon::ThreadPoolBuilder::new()
199            .num_threads(MIN_WORKERS_FOR_FORK_JOIN)
200            .build()
201            .unwrap()
202            .install(f)
203    }
204
205    #[test]
206    fn deep_chain_picks_level_scheduling() {
207        let pool = p();
208        let x = pool.symbol("x", Domain::Real);
209        let deep = junk(&pool, x, 300);
210        assert_eq!(
211            with_workers(|| choose_strategy(deep, &pool)),
212            Strategy::LevelScheduled
213        );
214    }
215
216    #[test]
217    fn wide_sum_picks_fork_join() {
218        let pool = p();
219        let zero = pool.integer(0_i32);
220        let args: Vec<ExprId> = (0..256)
221            .map(|i| {
222                let x = pool.symbol(format!("x{i}"), Domain::Real);
223                pool.add(vec![x, zero])
224            })
225            .collect();
226        let wide = pool.add(args);
227        assert_eq!(
228            with_workers(|| choose_strategy(wide, &pool)),
229            Strategy::ForkJoin
230        );
231    }
232
233    #[test]
234    fn few_workers_pick_level_scheduling() {
235        let pool = p();
236        let zero = pool.integer(0_i32);
237        let args: Vec<ExprId> = (0..256)
238            .map(|i| {
239                let x = pool.symbol(format!("y{i}"), Domain::Real);
240                pool.add(vec![x, zero])
241            })
242            .collect();
243        let wide = pool.add(args);
244        // Wide enough for fork-join on shape alone, but not enough workers.
245        let tp = rayon::ThreadPoolBuilder::new()
246            .num_threads(MIN_WORKERS_FOR_FORK_JOIN - 1)
247            .build()
248            .unwrap();
249        assert_eq!(
250            tp.install(|| choose_strategy(wide, &pool)),
251            Strategy::LevelScheduled
252        );
253    }
254
255    /// Whichever branch is taken, the answer must match the sequential engine.
256    #[test]
257    fn auto_matches_sequential_on_both_shapes() {
258        let pool = p();
259        let x = pool.symbol("x", Domain::Real);
260        let deep = junk(&pool, x, 200);
261        let zero = pool.integer(0_i32);
262        let args: Vec<ExprId> = (0..64)
263            .map(|i| {
264                let s = pool.symbol(format!("z{i}"), Domain::Real);
265                junk(&pool, s, 4)
266            })
267            .collect();
268        let wide = pool.add(args);
269        for expr in [deep, wide] {
270            let seq = simplify(expr, &pool).value;
271            let auto = with_workers(|| simplify_auto(expr, &pool).value);
272            assert_eq!(seq, auto);
273        }
274        let _ = zero;
275    }
276
277    #[test]
278    fn shape_measures_width_and_height() {
279        let pool = p();
280        let x = pool.symbol("x", Domain::Real);
281        // A chain of `Add(e, 0)` nodes: one node per level plus the leaves.
282        let deep = junk(&pool, x, 10);
283        let (nodes, height) = shape(deep, &pool);
284        assert!(
285            height >= 20,
286            "chain should be at least 20 levels, got {height}"
287        );
288        assert!(
289            (nodes as f64 / height as f64) < WIDTH_THRESHOLD,
290            "a chain must read as narrow"
291        );
292    }
293}
alkahest_cas/simplify/dispatch.rs

alkahest_cas/simplify/
dispatch.rs
