egraphs-good/egglog

Auto Demand Rewrite Conditions

saulshanabrook opened this issue · 5 comments

saulshanabrook commented

One thing I find myself often doing is having to "demand" the values I am conditioning against in rewrite rule. For example, you can see the manual version of this in the matrix example:

(rewrite (MMul (Kron A B) (Kron C D))
    (Kron (MMul A C) (MMul B D))
    :when
        ((= (ncols A) (nrows C))
        (= (ncols B) (nrows D)))
)

; demand
(rule ((= e (MMul A B)))
((ncols A)
(nrows A)
(ncols B)
(nrows B))
)

(rule ((= e (Kron A B)))
((ncols A)
(nrows A)
(ncols B)
(nrows B))
)

In order to match on the conditional, we first must add nrows and ncols to the graph.

I was thinking, what if we made this automatic when desugarring a rewrite with equality conditions? This could be done by adding another rule that matches on the left hand side of the rewrite and adds all expressions that are part of equality conditions in the actions of that rule.

For example, currently, the rewrite above desugars to something like:

(rule ((= x (MMul (Kron A B) (Kron C D)))
          (= (ncols A) (nrows C))
          (= (ncols B) (nrows D)))
         ((union x (Kron (MMul A C) (MMul B D)))

What if it also added this rule when desugarring:

(rule ((MMul (Kron A B) (Kron C D)))
      ((ncols A) (nrows C) (ncols B) (nrows D)))

Do you think this would ever be a problem to automatically introduce these extra demand rules? (I mean in examples we know about, not in the abstract). It could also be behind a rewrite flag to turn it on or off...

EDIT: Keeping track of every time this comes up

saulshanabrook commented

The fact that you have to manually demand rule conditions came up as a confusion on the Zulip recently: https://egraphs.zulipchat.com/#narrow/stream/375765-egg.2Fegglog/topic/Simple.20simplification.20with.20conditionals.2E/near/451613621

oflatt commented

I'm certain that auto-demanding every rewrite would cause performance problems in eggcc.
However, I'm not against introducing something like rewrite-and-demand that does this (perhaps just on the python side?).
I think we'll want to be careful since not every rewrite can be auto-demanded (i.e. if it depends on functions without defaults)

  • 👍1
ezrosent commented

We may want explicit syntax for some of this, see !/? in this paper

yihozhang commented

This proposal does not handle (rewrite e1 e2 :when ((= (get1 e1 c) (get2 e1 c))))

saulshanabrook commented

Here is a concrete example of @yihozhang's point:

(datatype Math
  (Num i64)
  (Var String)
  (Add Math Math)
  (Mul Math Math))

(rewrite (Var a) (Var b) :when ((= (Add (Var a) c) (Add (Var b) c))))

(union (Add (Var "a") (Num 0)) (Add (Var "b") (Num 0)))

(run 10)
(extract (Var "a") 10)

In this case we can't demand the conditions because c is unbound.