costfunction to prefer output as expanded polynomial

[nlewycky]

Hi! I'm having difficulty expressing a property of the EClass that I want to obtain through extraction.

I have a simple calculator with variables, u32's, Add and Mul, similar to the tests/math.rs example but simplified. I'd like to prefer output in the form of a sum of products, like a polynomial "x^2 + 3x + 4" or "(+ (* x x) (+ (* 3 x) 4))". The idea I have is to replace the CostFunction's Cost with an object that tracks "is_add" and "is_mul" to report a low cost for add-of-add and add-of-mul and mul-of-mul, but high costs otherwise.

#[derive(Default, Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Hash)]
pub struct PolyCost {
    cost: u32,
    is_add: bool,
    is_mul: bool,
}
impl PolyCost {
    fn merge_child_into(&mut self, child: Self) -> Self {
        if self.is_add && child.is_add {
            self.cost += child.cost + 1;
        } else if child.is_mul && self.is_add {
            self.cost += child.cost + 1;
        } else if child.is_mul && self.is_mul {
            self.cost += child.cost + 1;
        } else if child.is_add {
            self.cost += child.cost + 100;
        } else if child.is_mul {
            self.cost += child.cost + 10000;
        } else {
            self.cost += child.cost + 1;
        }
        *self
    }
    fn get_cost(&self) -> u32 {
        if self.is_add {
            self.cost + 100
        } else if self.is_mul {
            self.cost + 10000
        } else {
            self.cost + 1
        }
    }
}
impl Ord for PolyCost {
    fn cmp(&self, other: &Self) -> Ordering {
        self.get_cost().cmp(&other.get_cost())
    }
}
impl std::fmt::Display for PolyCost {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        write!(fmt, "{}", self.get_cost())
    }
}

pub struct ConstantFold {}
impl CostFunction<Math> for ConstantFold {
    type Cost = PolyCost;

    fn cost<C>(&mut self, enode: &Math, mut costs: C) -> Self::Cost
    where
        C: FnMut(Id) -> Self::Cost,
    {
        match enode {
            Math::Add(..) => {
                let p = PolyCost {
                    is_add: true,
                    ..Default::default()
                };
                enode.fold(p, |mut sum, i| sum.merge_child_into(costs(i)))
            }
            Math::Mul(..) => {
                let p = PolyCost {
                    is_mul: true,
                    ..Default::default()
                };
                enode.fold(p, |mut sum, i| sum.merge_child_into(costs(i)))
            }
            _ => {
                let p = PolyCost {
                    cost: 1,
                    ..Default::default()
                };
                enode.fold(p, |mut sum, i| sum.merge_child_into(costs(i)))
            }
        }
    }
}

fn make_rules() -> Vec<Rewrite<Math, ConstantFold>> {
    vec![
        rewrite!("add-0"; "(+ 0 ?a)" => "?a"),
        rewrite!("mul-0"; "(* 0 ?a)" => "0"),
        rewrite!("mul-1"; "(* 1 ?a)" => "?a"),
        rewrite!("commute-add"; "(+ ?a ?b)" => "(+ ?b ?a)"),
        rewrite!("commute-mul"; "(* ?a ?b)" => "(* ?b ?a)"),
        rewrite!("assoc-add"; "(+ ?a (+ ?b ?c))" => "(+ (+ ?a ?b) ?c)"),
        rewrite!("assoc-mul"; "(* ?a (* ?b ?c))" => "(* (* ?a ?b) ?c)"),
        rewrite!("distribute"; "(* ?a (+ ?b ?c))" => "(+ (* ?a ?b) (* ?a ?c))"),
        rewrite!("factor"; "(+ (* ?a ?b) (* ?a ?c))" => "(* ?a (+ ?b ?c))"),
        rewrite!("add-self"; "(+ ?a ?a)" => "(* 2 ?a)"),
        rewrite!("rev-add-self"; "(* 2 ?a)" => "(+ ?a ?a)"),
    ]
}

Unfortunately this doesn't do the trick:

$ ./target/debug/x '(* (* (+ x 1) (+ x 2)) (* (+ x 3) (+ x 4)))'
Simplified (* (* (+ x 1) (+ x 2)) (* (+ x 3) (+ x 4))) to (* (+ (+ (* x 4) (* x x)) 3) (+ (+ (+ 8 x) (+ x (* x x))) (* x 4))) with cost 10232

I have a theory that this might be because of instability in the ordering of the costs, a cost might appear to be more expensive until it's merged into the enode parent. I wasn't able to find documentation on what the requirements on Cost are, to suggest whether this sort of "negative cost" is permitted. I'd appreciate understanding why this fails, or better, how to do it.

[mwillsey]

I think the issue is that this is ultimately a non-local cost function. In this case, you might true collapsing the pattern you're looking for into a single e-node. Perhaps with a add-multiply (fma) operator that has very low cost and a rewrite like:

(+ (* a b) c) -> (fma a b c)

[Bastacyclop]

Without looking into the details, my intuition would be that your Ord implementation is greedily prefering sub-expression with lower cost, even though using a sub-expression with higher cost might be more optimal once put into a surrounding expression context. If that it the case, #251 (comment) should help.

[nlewycky]

Unless I didn't find it, there isn't documentation for the rules on Cost, besides that it must implement PartialOrd (but really Ord or else egg will panic!). Are negative costs permitted? If costs a > b then must a+x > b+x? Or are such rules up to the particulars of the implementation of the Extractor? I'd hope not because it seems to be an important feature connected to how the e-graph encodes an infinite graph specifically so that it can explore all possibilities, but regardless of what the rules are I'd like to have them documented.

Beyond that, if the answer turns out to be that I need a custom Extractor, I can live with that.

[nlewycky]

Are there sub-expressions where the locally best term is different from the globally best term (i.e. the best choice of sub-expression depends on the surrounding expression)?

Yes. The approach is to make trees-of-adds cheap, trees-of-multiplies cheap and multiplies that are children of adds cheap, everything else is expensive.

[mwillsey]

I basically think that the current extractor does not support what you're looking for. You are right, we could do better to document the "monotonicity" requirements of the current extractor.

[nlewycky]

FWIW, I found that there are already doc-comments which explain this requirement that I previously missed:

egg/src/extract.rs

Lines 113 to 114 in 164d2b3

	You might want to use [`saturating_add`](u64::saturating_add) to
	ensure your cost function is still monotonic in this situation.

and

egg/src/extract.rs

Lines 49 to 51 in 164d2b3

	To extract an expression from an [`EGraph`], the [`Extractor`]
	requires a cost function to performs its greedy search.
	`egg` provides the simple [`AstSize`] and [`AstDepth`] cost functions.

[Kixiron]

Are there sub-expressions where the locally best term is different from the globally best term (i.e. the best choice of sub-expression depends on the surrounding expression)?

There's plenty of situations where a "side effectful" (performance wise) function before (above, in the egraph) it could change costs, e.g. the (index P) of (join (index P) (index Q)) could be free when indexes can be reused once made, making it more costly when another (index P) exists
Another example could be some similar kind of allocation or effort reuse, like reusing previously allocated buffers.

[Bastacyclop]

Alternatively, you might want to keep track of multiple best candidates per e-class to choose the right one depending on context?

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub enum Context {
  IsAdd,
  IsMul,
  IsOther
}

#[derive(Default, Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Hash)]
pub struct PolyCost {
    costs: HashMap<Context, u32>,
}

// [...]

You can't use the vanilla extractor with this, because you need to associate an L node for each cost (in reality, costs: HashMap<Context, (u32, L)>), but you can achieve this type of extraction using an e-class analysis.

[Bastacyclop]

I've used this in the past to implement various extraction strategies (e.g. #174), and would love to explore more possibilities in this direction.