Rewrite raw egglog code in Luminal2 into Rust API by Eggplant crate

[MilkBlock]

Recently I opened a crate called eggplant which is a High-Level Rust API for egglog providing easiest approach to define a suite of DSL and Rules (Patttern Action).

With eggplant you can define a DSL and implement constant propagation in only 50 lines of code in Rust! Ref

You can get great flexibility when coding an action in Rust comparing to raw egglog code

I'm willing to rewrite current raw egglog code in Luminal2 into Rust API through this crate. As first step, I will implement FlashAttention Rewrite as a demo.

Could you have a look and give me some suggestion? @jafioti

[jafioti]

Wow very cool! Are there any major limitations with your API vs directly using egglog? Are there perf improvements? Currently we render the dags directly to egglog code as a string and feed it into egglog, but I'd love a more "native" approach.

[MilkBlock]

Wow very cool! Are there any major limitations with your API vs directly using egglog? Are there perf improvements? Currently we render the dags directly to egglog code as a string and feed it into egglog, but I'd love a more "native" approach.

As for Action, eggplant directly exposes egglog-backend's Execution State as closure arguemnt. That means you can render dags through this struct with zero-cost. The crate compensate the drawback of Execution State which receives a unknown-length list of untyped Values (error prone pure ID with no type info).

As for Pattern, it is more elegant then egglog code because

1. It allows you to specify discriminant, and you can get basic components(i64 i32 String ...) of this discriminant.
   This is the core idea of eggplant without which it's hard to specify patterns variables because of the parallax between basic component and complex components(fields that are not basic component).

#[eggplant::pat_vars]
struct MulCommuVars {
    l_expr: Expr,  // discriminant unspecified
    r_expr: Expr,  // you can't get basic components because discriminant is unspecified
    p_expr: Mul,  // mul have no basic components 
}

#[eggplant::pat_vars]
struct MulPropVars {
    l_expr: Const,  // discriminant specified
    r_expr: Const,  // now you can get  r_expr.num in Action context, which keeps full type info.
    p_expr: Mul,
}

while as for egglog you have to directly declare the name of basic components in rule.

For example:

(rewrite (gcon I a) a)

In fact, we don't need to specify the name of a.

Lackness:

Pattern Predicate API is on schedule, now some check has to be delayed in Action.

[MilkBlock]

It also supports directly send egglog code into inner EGraph. That means it's possible to incrementally replace current implementation into Rust Native Impl.

[jafioti]

@MilkBlock Yeah this looks cool, i'd be super interested in what a small luminal_2 reproduction would look like to see if it reduces complexity

[jafioti]

@MilkBlock Do you know how to get explanations out of egglog? I would love to be able to pass in two expressions and get a trajectory of rewrites connecting them. I remember egg did this, but I can't figure out how to do it in egglog.

[MilkBlock]

@MilkBlock Do you know how to get explanations out of egglog? I would love to be able to pass in two expressions and get a trajectory of rewrites connecting them. I remember egg did this, but I can't figure out how to do it in egglog.

@jafioti
https://github.com/MilkBlock/egglog/blob/eggplant_test_use/examples/proof_test.rs
I wrote an example on this. Ping me if you have any question.

This trick requires to do some revision on egglog (enable with_tracing backend and pub egraph backend field)
Remember don't use (let xxxx xxx) grammar. Use egglog::prelude::expr to create node, or if you want to create node in text, define a rule of nodes creation and run it once.

[jafioti]

@MilkBlock This is awesome!

I tried to modify it to get string vars created because I couldn't get (MVar "a") working in the expr! macro:

use egglog::{ast::Expr, expr, fact, span, EGraph, SerializeConfig, TermDag, Value};

fn main() {
    let mut egraph = egglog::EGraph::with_tracing();
    declare(&mut egraph);

    // TermDag::
    // use egglog::prelude::exprs;
    // let value1 = expr!((MAdd (MNum 0) (MNum 2)));
    // let value2 = expr!((MNum 2));
    // let value1 = egraph.eval_expr(&value1).unwrap().1;
    // let value2 = egraph.eval_expr(&value2).unwrap().1;

    egraph
        .parse_and_run_program(
            None,
            "
(let a (MAdd (MNum 0) (MNum 2)))
(let b (MNum 2))
    ",
        )
        .unwrap();

    let (_, value1) = egraph.eval_expr(&egglog::var!("a")).unwrap();
    let (_, value2) = egraph.eval_expr(&egglog::var!("b")).unwrap();

    egraph
        .parse_and_run_program(
            None,
            "
(let a (MAdd (MNum 0) (MNum 2)))
(let b (MNum 2))
        (run 10)
    ",
        )
        .unwrap();
    let proof = &egraph
        .backend
        .explain_terms_equal(value1, value2)
        .unwrap()
        .reason;
    println!("{}", proof);
}

But I ran into proofs aren't supported for non-union merge functions
Is this expected?

[MilkBlock]

@jafioti
Seems they missed non-int literal constructor on exprs macro. An alternative method is to build this TermDag yourself,

    let mut term_dag = TermDag::default();
    let a = term_dag.lit(Literal::String("a".to_string()));
    let b = term_dag.lit(Literal::String("b".to_string()));

    let (_, value1) = egraph
        .eval_expr(&term_dag.term_to_expr(&a, span!()))
        .unwrap();
    let (_, value2) = egraph
        .eval_expr(&term_dag.term_to_expr(&b, span!()))
        .unwrap();

Anyway, I think this is kind of wordy. Soon I will give an Proof demo by Eggplant to help Debug.

[MilkBlock]

@jafioti
Revision.
A better way is to use egraph.parser.get_expr_from_string

    let expr1 = egraph
        .parser
        .get_expr_from_string(None, "(MAdd (MNum 2) (MNum 0))")
        .unwrap();
    let expr2 = egraph
        .parser
        .get_expr_from_string(None, "(MNum 2)")
        .unwrap();
    let (_, v1) = egraph.eval_expr(&expr1).unwrap();
    let (_, v2) = egraph.eval_expr(&expr2).unwrap();

[jafioti]

@MilkBlock Awesome, this looks great. Do you plan to upstream all of this into egglog, or a subset and then build eggplant off of it? Currently we try to use the latest stable version of egglog on their github. Are you planning to merge this with them soon? would love to adopt it, especially the explanations.

[MilkBlock]

I will be glad if egglog expose the explanations API but it depends on egglog team's progress. In fact, don't think direct API call is the final interface for users, it should be wrapped as UI and we may only click on two e-node in egraph to get the explanation.

So I decided to do explanation UI feature for eggplant and maintain a fork of egglog.

[jafioti]

@MilkBlock I'm adding in your egglog fork to luminal to get proofs. Is there a way to define multi-line expressions, such as

(let a (MNum 0))
(let b (MNum 1))
(let c (MAdd a b))
(let d (MNum 1))

and get a proof that c == d?
currently it seems like get_expr_from_string and eval_expr need one-line expressions

[MilkBlock]

In egglog, a Symbol(a,b,c,d) is represented as Function Table, but currently egglog-backend or egglog doesn't support function-related proof generation(only regular constructors are allowed), so this might take lots of effort.

The best way to enable proof is avoid Function Table representation. So eggplant directly uses egglog::Value to represent expression rather than Symbol.

Given this dilemma, all I can do is enabling debug feature in eggplant as soon as possible. I think maintaining a copy of DSL and rules in eggplant is not hard compared to revising egglog.

So I recommend debugging with eggplant in test branch but keep main branch in egglog which both keep stability and enable debugging.

I think I can give a demo in egglog with all equivalent rules in Luminal FlashAttentionDemo after few days with basic debugging(without UI).

[jafioti]

awesome, let me know if i can help. it would be incredible to be able to have 2 different luminal dags, a base one and a transformed one, convert them both into egglog, and generate a proof showing how the rewrite rules connect the two.

again let me know if i can help!