Hi from the Julia lang planet! 14.5 times faster than egg!

[0x0f0f0f]

I'm really happy to announce that after some years of work, and hardcore optimizations, I've made Metatheory.jl 14.5 times faster than egg in rust! (at least on a simple example)

This is the experimental branch JuliaSymbolics/Metatheory.jl#185

We should definitely work together on something!! I'm really excited to work on porting many applications, such as Herbie, now that performance is on-par with egg!

This implementation works on the Julia AST direclty, and allows Julia code to be rewritten with e-graphs before compilation (in hygienic macros, as in scheme and other LISPs).

Other than rewriting just the surface Julia AST, this could be really easily plugged into the compiler, rewriting directly on the lowered typed SSA form, in order to automatically discover new compiler optimizations!!

Julia also has a lot of packages that rewrite numerical expressions symbolically before evaluating them numerically, and this would be a super nice research direction, as we already showcased in an example in https://dl.acm.org/doi/pdf/10.1145/3511528.3511535

Julia Code

using Metatheory, BenchmarkTools

t = @theory a b begin
  a + b --> b + a
  a * b --> b * a
  a + 0 --> a
  a * 0 --> 0
  a * 1 --> a
end

using BenchmarkTools

p = SaturationParams(; timer = false)

function simpl(ex)
  g = EGraph(ex)
  saturate!(g, t, p)
  extract!(g, astsize)
end

ex = :(0 + (1 * foo) * 0 + (a * 0) + a)

simpl(ex)

@btime simpl(ex)

94.462 μs (1412 allocations: 66.08 KiB)

Rust code

use egg::{rewrite as rw, *};
//use std::time::Duration;
fn main() {
    env_logger::init();
    use egg::*;

    define_language! {
        enum SimpleLanguage {
            Num(i32),
            "+" = Add([Id; 2]),
            "*" = Mul([Id; 2]),
            Symbol(Symbol),
        }
    }

    fn make_rules() -> Vec<Rewrite<SimpleLanguage, ()>> {
        vec![
            rewrite!("commute-add"; "(+ ?a ?b)" => "(+ ?b ?a)"),
            rewrite!("commute-mul"; "(* ?a ?b)" => "(* ?b ?a)"),
            rewrite!("add-0"; "(+ ?a 0)" => "?a"),
            rewrite!("mul-0"; "(* ?a 0)" => "0"),
            rewrite!("mul-1"; "(* ?a 1)" => "?a"),
        ]
    }

    /// parse an expression, simplify it using egg, and pretty print it back out
    fn simplify(s: &str) -> String {
        // parse the expression, the type annotation tells it which Language to use
        let expr: RecExpr<SimpleLanguage> = s.parse().unwrap();

        // simplify the expression using a Runner, which creates an e-graph with
        // the given expression and runs the given rules over it
        let runner = Runner::default().with_expr(&expr).run(&make_rules());

        // the Runner knows which e-class the expression given with `with_expr` is in
        let root = runner.roots[0];

        // use an Extractor to pick the best element of the root eclass
        let extractor = Extractor::new(&runner.egraph, AstSize);
        let (best_cost, best) = extractor.find_best(root);
        println!("Simplified {} to {} with cost {}", expr, best, best_cost);
        best.to_string()
    }

    // assert_eq!(simplify("(* 0 42)"), "0");
    let apply_time: std::time::Instant = instant::Instant::now();
    // assert_eq!(simplify("(+ 0 (* 1 foo))"), "foo");
    assert_eq!(simplify("(+ (+ (+ 0 (* (* 1 foo) 0)) (* a 0)) a)"), "a");
    let apply_time = apply_time.elapsed().as_secs_f64();
    println!("simplification time {}", apply_time);
}

simplification time 0.001375786 seconds which is 1375.786 microseconds

1375.786 / 94.462 = 14.56 :)

[mwillsey]

Congrats! That is awesome!! Very exciting to see that the Julia community is (still!) excited about these techniques.

You might also consider sharing this on the EGRAPHS zulip, as a lot of the discussion has moved there. I'd be interested in hearing about what techniques you feel are the most important to the performance improvements.