LpExtractor creates an infeasible problem for CBC

[khaki3]

Hi. I'm having fun with egg.

I encountered an infeasible error with CBC that prevents LpExtractor from processing cyclic e-graphs.

Example:

use egg::*;

fn main() {
    env_logger::init();

    let mut egraph = EGraph::<SymbolLang, ()>::default();

    let g = egraph.add_expr(&"(g v x)".parse().unwrap());
    let v = egraph.add_expr(&"v".parse().unwrap());
    let f = egraph.add(SymbolLang::new("f", vec![v, g]));
    let top = egraph.add(SymbolLang::new("list", vec![f, g]));
    egraph.rebuild();

    let runner = Runner::default()
        .with_iter_limit(100)
        .with_node_limit(10_000)
        .with_egraph(egraph)
        .run(&vec![rewrite!("f"; "(f ?v (g ?v ?a))" => "?a")]);

    runner.egraph.dot().to_dot("tmp.dot").unwrap();

    LpExtractor::new(&runner.egraph, AstSize).solve(top);
}

Log produced:

[INFO  egg::egraph] REBUILT! in 0.000s
      Old: hc size 5, eclasses: 5         
      New: hc size 5, eclasses: 5
      unions: 0, trimmed nodes: 0
[INFO  egg::egraph] REBUILT! in 0.000s
      Old: hc size 5, eclasses: 5        
      New: hc size 5, eclasses: 5
      unions: 0, trimmed nodes: 0
[INFO  egg::run]
    Iteration 0
[INFO  egg::run] Search time: 0.000048125
[INFO  egg::run] Apply time: 0.000038416
[INFO  egg::egraph] REBUILT! in 0.000s
      Old: hc size 5, eclasses: 4
      New: hc size 6, eclasses: 4
      unions: 0, trimmed nodes: 0
[INFO  egg::run] Rebuild time: 0.000590618
[INFO  egg::run] Size: n=6, e=4
[INFO  egg::run]
    Iteration 1
[INFO  egg::run] Search time: 0.000033264
[INFO  egg::run] Apply time: 0.000013224
[INFO  egg::egraph] REBUILT! in 0.000s
      Old: hc size 6, eclasses: 4
      New: hc size 6, eclasses: 4
      unions: 0, trimmed nodes: 0
[INFO  egg::run] Rebuild time: 0.000198928
[INFO  egg::run] Size: n=6, e=4
[INFO  egg::run] Stopping: Saturated
[/home/user/.cargo/registry/src/github.com-1ecc6299db9ec823/egg-0.9.0/src/lp_extract.rs:138] max_order = 50.0
Welcome to the CBC MILP Solver
Version: 2.10
Build Date: Jun 30 2022

command line - Cbc_C_Interface -solve -quit (default strategy 1)
Problem is infeasible - 0.00 seconds
Total time (CPU seconds):       0.01   (Wallclock seconds):       0.00

[INFO  egg::lp_extract] CBC status Finished, LinearRelaxationInfeasible
thread 'main' panicked at 'called `Option::unwrap()` on a `None` value', /home/user/.cargo/registry/src/github.com-1ecc6299db9ec823/egg-0.9.0/src/lp_extract.rs:192:80
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

This find_cycles extracts g as a result. But I might need f instead.
https://github.com/egraphs-good/egg/blob/main/src/lp_extract.rs#L209

For the moment, I'm avoiding the error with code like is_dag. Perhaps, it works for others.

fn find_cycles<L, N>(egraph: &EGraph<L, N>, mut f: impl FnMut(Id, usize))
where
    L: Language,
    N: Analysis<L>,
{
    for class in egraph.classes() {
        let id = class.id;
        for (i, node) in egraph[id].iter().enumerate() {
            for &child in node.children() {
                if usize::from(child) >= usize::from(id) {
                    f(id, i);
                    break;
                }
            }
        }
    }
}

(Could be related: #196 )

[khaki3]

The code above does not support additional nodes. I ended up with this unpolished one:

fn find_cycles<L, N>(egraph: &EGraph<L, N>, mut f: impl FnMut(Id, usize))
where
    L: Language,
    N: Analysis<L>,
{
    let mut pending: HashMap<Id, Vec<(Id, usize)>> = HashMap::default();

    let mut order: HashMap<Id, usize> = HashMap::default();

    let mut memo: HashMap<(Id, usize), bool> = HashMap::default();

    let mut stack: Vec<(Id, usize)> = vec![];

    for class in egraph.classes() {
        let id = class.id;
        for (i, node) in egraph[id].iter().enumerate() {
            for &child in node.children() {
                pending.entry(child).or_insert_with(Vec::new).push((id, i));
            }

            if node.is_leaf() {
                stack.push((id, i));
            }
        }
    }

    let mut count = 0;

    while let Some((id, i)) = stack.pop() {
        if memo.get(&(id, i)).is_some() {
            continue;
        }

        let node = &egraph[id].nodes[i];
        let mut update = false;

        if node.is_leaf() {
            update = true;
        } else if node.children().iter().all(|&x| order.get(&x).is_some()) {
            if let Some(ord) = order.get(&id) {
                update = node.children().iter().all(|&x| order.get(&x).unwrap() < ord);
                if !update {
                    memo.insert((id, i), false);
                    continue;
                }
            } else {
                update = true;
            }
        }

        if update {
            if order.get(&id).is_none() {
                order.insert(id, count);
                count = count + 1;
            }
            memo.insert((id, i), true);
            if let Some(mut v) = pending.remove(&id) {
                stack.append(&mut v);
                stack.sort();
                stack.dedup();
            };
        }
    }

    for class in egraph.classes() {
        let id = class.id;
        for (i, node) in egraph[id].iter().enumerate() {
            if let Some(true) = memo.get(&(id, i)) {
                continue;
            }
            f(id, i);
        }
    }
}

[mwillsey]

Hi. I'm having fun with egg.

Great to hear! Thanks for providing a workaround that seems to work here. I'm not quite sure I fully understand what is deficient in the existing find_cycles. Do you?

[khaki3]

The current one starts the search from the top of classes() and preserves them in its order. For allowing my case, it would need to construct a RecExpr-like order. Expressions fed into an egraph originally have such orders and recovering them is partially possible while some of them become cycles during rewrites (e.g. f is unioned with x in my example) but always have acyclic nodes in the same eclass so can be ignored for the extraction.

[khaki3]

My code is missing some optimal cases. The first one could be supported by checking actual cycles (not existing in this case so no remove) after constructing a RecExpr-like order. But the second one having interdependent cycles would need a cost function to choose but it's an expensive search.

use egg::*;

struct CostFn;

type EGraph = egg::EGraph<SymbolLang, ()>;

impl LpCostFunction<SymbolLang, ()> for CostFn {
    fn node_cost(&mut self, _egraph: &EGraph, _eclass: Id, enode: &SymbolLang) -> f64
    {
        match enode.op.as_str() {
            "A" => 1.0,
            "B" => 2.0,
            "x" => 10.0,
            "y" => 3.0,
            _ => 0.0,
        }
    }
}

fn main() {
    env_logger::init();

    let mut egraph = EGraph::default();

    let a = egraph.add_expr(&"(A y)".parse().unwrap());
    let x = egraph.add_expr(&"x".parse().unwrap());
    let y = egraph.add_expr(&"y".parse().unwrap());
    let top = egraph.add(SymbolLang::new("list", vec![a, y]));
    egraph.union(a, x);
    egraph.rebuild();
    egraph.dot().to_dot("tmp1.dot").unwrap();
    println!("{}", LpExtractor::new(&egraph, CostFn).solve(top)); // (list x y) while (list (A y) y) is optimal 

    let b = egraph.add_expr(&"(B x)".parse().unwrap());
    egraph.union(b, y);
    egraph.rebuild();
    egraph.dot().to_dot("tmp2.dot").unwrap();
    println!("{}", LpExtractor::new(&egraph, CostFn).solve(top)); // (list x (B x)) while (list (A y) y) is optimal
}

[mwillsey]

Hmm, this seems tricky indeed. @Bastacyclop do you think this approach could solve #196?

[Bastacyclop]

Hi @mwillsey , I'll investigate next week!

[Bastacyclop]

Using @khaki3 's find_cycles on one of my problematic benchmark fixes the issue (at least a solution is found). I'll use this change in more benchmarks and report if it holds up.

[Bastacyclop]

Using this solution on a broader set of benchmarks produced mitigated results for us: extraction times are increased (leading to timeouts), and extracted solutions are not always better. I'm currently attempting to implement an alternative solution, if you have ideas on benchmarks to evaluate on.