Implement custom e-matching benchmarking

.gitignore

@@ -13,4 +13,5 @@ default.nix
 TODO.md
 flamegraph.svg
 perf.data*
+*.bench
 .vscode/settings.json

Cargo.toml

@@ -28,20 +28,10 @@ vectorize = { version = "0.2", optional = true }
 # for the reports feature
 serde_json = { version = "1", optional = true }
 
-[[bench]]
-name = "bench_tests"
-harness = false
-
 [dev-dependencies]
 env_logger = {version = "0.7", default-features = false}
 ordered-float = "1"
 
-[dev-dependencies.iai]
-# version = "*"
-git = "https://github.com/mwillsey/iai"
-rev = "c0b7cb47181e9577a6847caae53a2f0f881f45b2"
-features = ["threadpool"]
-
 [features]
 wasm-bindgen = [ "instant/wasm-bindgen" ]
 serde-1 = [ "serde", "indexmap/serde-1", "hashbrown/serde", "vectorize" ]

src/test.rs

@@ -103,6 +103,94 @@ pub fn test_runner<L, A>(
     }
 }
 
+fn percentile(k: f64, data: &[u128]) -> u128 {
+    // assumes data is sorted
+    assert!((0.0..=1.0).contains(&k));
+    let i = (data.len() as f64 * k) as usize;
+    let i = i.min(data.len() - 1);
+    data[i]
+}
+
+pub fn bench_egraph<L, N>(_name: &str, rules: Vec<Rewrite<L, N>>, exprs: &[&str]) -> EGraph<L, N>
+where
+    L: Language + FromOp + 'static + Display,
+    N: Analysis<L> + Default + 'static,
+{
+    let mut patterns: Vec<Pattern<L>> = vec![];
+    for rule in &rules {
+        if let Some(ast) = rule.searcher.get_pattern_ast() {
+            patterns.push(ast.alpha_rename().into())
+        }
+        if let Some(ast) = rule.applier.get_pattern_ast() {
+            patterns.push(ast.alpha_rename().into())
+        }
+    }
+    eprintln!("{} patterns", patterns.len());
+
+    patterns.retain(|p| p.ast.as_ref().len() > 1);
+    patterns.sort_by_key(|p| p.to_string());
+    patterns.dedup();
+    patterns.sort_by_key(|p| p.ast.as_ref().len());
+
+    let iter_limit = env_var("EGG_ITER_LIMIT").unwrap_or(1);
+    let node_limit = env_var("EGG_NODE_LIMIT").unwrap_or(1_000_000);
+    let time_limit = env_var("EGG_TIME_LIMIT").unwrap_or(1000);
+    let n_samples = env_var("EGG_SAMPLES").unwrap_or(100);
+    eprintln!("Benching {} samples", n_samples);
+    eprintln!(
+        "Limits: {} iters, {} nodes, {} seconds",
+        iter_limit, node_limit, time_limit
+    );
+
+    let mut runner = Runner::default()
+        .with_scheduler(SimpleScheduler)
+        .with_hook(move |runner| {
+            let n_nodes = runner.egraph.total_number_of_nodes();
+            eprintln!("Iter {}, {} nodes", runner.iterations.len(), n_nodes);
+            if n_nodes > node_limit {
+                Err("Bench stopped".into())
+            } else {
+                Ok(())
+            }
+        })
+        .with_iter_limit(iter_limit)
+        .with_node_limit(node_limit)
+        .with_time_limit(Duration::from_secs(time_limit));
+
+    for expr in exprs {
+        runner = runner.with_expr(&expr.parse().unwrap());
+    }
+
+    let runner = runner.run(&rules);
+    eprintln!("{}", runner.report());
+    let egraph = runner.egraph;
+
+    let get_len = |pat: &Pattern<L>| pat.to_string().len();
+    let max_width = patterns.iter().map(get_len).max().unwrap_or(0);
+    for pat in &patterns {
+        let mut times: Vec<u128> = (0..n_samples)
+            .map(|_| {
+                let start = Instant::now();
+                let matches = pat.search(&egraph);
+                let time = start.elapsed();
+                let _n_results = matches.iter().map(|m| m.substs.len()).sum::<usize>();
+                time.as_nanos()
+            })
+            .collect();
+        times.sort_unstable();
+
+        println!(
+            "test {name:<width$} ... bench: {time:>10} ns/iter (+/- {iqr})",
+            name = pat.to_string().replace(' ', "_"),
+            width = max_width,
+            time = percentile(0.05, &times),
+            iqr = percentile(0.75, &times) - percentile(0.25, &times),
+        );
+    }
+
+    egraph
+}
+
 /// Make a test function
 #[macro_export]
 macro_rules! test_fn {

tests/lambda.rs

@@ -339,3 +339,40 @@ egg::test_fn! {
         (app (var fib) 4))"
     => "3"
 }
+
+#[test]
+fn lambda_ematching_bench() {
+    let exprs = &[
+        "(let zeroone (lam x
+            (if (= (var x) 0)
+                0
+                1))
+            (+ (app (var zeroone) 0)
+            (app (var zeroone) 10)))",
+        "(let compose (lam f (lam g (lam x (app (var f)
+                                        (app (var g) (var x))))))
+        (let repeat (fix repeat (lam fun (lam n
+            (if (= (var n) 0)
+                (lam i (var i))
+                (app (app (var compose) (var fun))
+                    (app (app (var repeat)
+                            (var fun))
+                        (+ (var n) -1)))))))
+        (let add1 (lam y (+ (var y) 1))
+        (app (app (var repeat)
+                (var add1))
+            2))))",
+        "(let fib (fix fib (lam n
+            (if (= (var n) 0)
+                0
+            (if (= (var n) 1)
+                1
+            (+ (app (var fib)
+                    (+ (var n) -1))
+                (app (var fib)
+                    (+ (var n) -2)))))))
+            (app (var fib) 4))",
+    ];
+
+    egg::test::bench_egraph("lambda", rules(), exprs);
+}

tests/math.rs

@@ -328,3 +328,18 @@ fn assoc_mul_saturates() {
 
     assert!(matches!(runner.stop_reason, Some(StopReason::Saturated)));
 }
+
+#[test]
+fn math_ematching_bench() {
+    let exprs = &[
+        "(i (ln x) x)",
+        "(i (+ x (cos x)) x)",
+        "(i (* (cos x) x) x)",
+        "(d x (+ 1 (* 2 x)))",
+        "(d x (- (pow x 3) (* 7 (pow x 2))))",
+        "(+ (* y (+ x y)) (- (+ x 2) (+ x x)))",
+        "(/ 1 (- (/ (+ 1 (sqrt five)) 2) (/ (- 1 (sqrt five)) 2)))",
+    ];
+
+    egg::test::bench_egraph("math", rules(), exprs);
+}