LIN-graphs, an Egraph modulo T

A LIN-graph has terms which are linear combinations of other terms. You could use this to model actual linear equations, or you could use it to model multiplication and division, or use it to model multi-sets (assuming you have interpretation of negative counts), or a number of other use cases. Linear combinations are associative and commutative, so this is basically a solution to AC matching, assuming you're OK with negative quantities.

Let's step through the recipe book from the last post. We have our language of terms: linear combinations of EUF terms, which we know we can think of as e-classes. We normalize linear combinations by sorting and combining like terms; we can sort in e-class ID order. Note that we'll need to sort and combine any time two e-classes are union'd.

Next we need a language of patterns. These have to include at least terms and pattern variables; I'm going to choose linear combinations of pattern variables and e-class IDs. For example, you might try matching sin(2 ?x), where 2 ?x is a linear pattern. Or log(?x) + log(?y); here the addition is a linear pattern. Recall that these patterns flatten into individual constriants, like so:

?z = log(?x)
?w = log(?y)
?r = ?z + ?w

Now that we have patterns, we need to define a set of cells, which include at least ?var = term but can potentially be richer. In our case, we won't need anything besides the basics, so our cells will just be bindings of pattern variables to terms. Patterns need to be closed under substitution of cells. This is why I defined patterns to contain both pattern variables and e-class IDs: so that substituting a term (containing e-class IDs) for a pattern variable (in a pattern) yields a pattern. Cells need to be closed under conjunction, but in this simple case we don't actually need to conjoin anything so this part is trivial.

Now let's discuss pattern matching. In pattern matching, we are given a constraint, asked to quantifier-eliminate all other variables from it. This is generally done by each theory separately for its constraints. Let's see that in action for normal e-graphs first. Suppose we get the constraint:

?z = log(?x)

Now, in this scenario, ?z is an e-class ID, as is (let's say) ?x. Now, we pattern match one variable at a time; say we are pattern-matching ?x. In that case we're trying to convert this logical formula:

∃ ?z, ?z = log(?x)

into a disjunction of cells. We do that by looking up all log nodes in our e-graph and collecting all of their arguments to form a set of ?x values. That gives us cells of the form ?x = <eclass ID>. Then we can substitute that in, so we get the new pattern ?z = log(e1); when we convert this to a cell we get an e-class ID for ?z.

Now let's look at the same for LIN-patterns. We have the equation:

?r = ?z + ?w

Let's say we're looking for ?w values. So, in other words, we're trying to quantifier-eliminate this equation:

∃ ?r ?w, ?r = ?z + ?w

Well, for what ?w do such variables exist? All of them, actually, so this equation yields no constraints on ?w. That's OK, actually, because other constraints on ?w will yield values for it.

But there are other cases. For example, suppose we first pattern-match and substitute the ?z variable for e-class e1 and the ?w variable for e-class e4. Then we have the constraint ?r = e1 + e4, which we can now solve for ?r, yielding, well, the obvious thing. Then that can be substituted into all other constraints.

Here's a complex pattern where you can see that happen:

f(?x + ?y) + g(?x - ?y)

Flattened, we get this:

?a = ?x + ?y
?b = f(?a)
?c = ?x - ?y
?d = g(?c)
?e = ?b + ?d

We can solve this following any variable order, but let's go in order: ?a, ?b, ?c, ?d, ?e, ?x, ?y, ?z. First we solve for ?a; that participates in two constraints. The first constraint, ∃ ?x ?y, ?a = ?x + ?y, is trivial, so doesn't constrain ?a in any way. The second constraint, ∃ ?b, ?b = f(?a), constrains ?a to be an argument to f; there might be more than one of these, but let's say we chose ?a = e1. Now we have a new query:

e1 = ?x + ?y
?b = f(e1)
?c = ?x - ?y
?d = g(?c)
?e = ?b + ?d

Let's solve for ?b; the first constraint, ?b = f(e1), yields some e-class ID for ?b, let's say e2. The last constraint, ∃ ?e ?d, ?e = ?b + ?d, again yields nothing, which is just as well. We now have this pattern:

e1 = ?x + ?y
e2 = f(e1)
?c = ?x - ?y
?d = g(?c)
?e = e2 + ?d

The second constraints, e2 = f(e1), is now fully substituted, so we can evaluate it (it's true) and discard it. You can image how ?c and ?d go; we're now at:

e1 = ?x + ?y
e3 = ?x - ?y
?e = e2 + e4

Next, we solve for ?e; this is easy because the final constraint just gives us its value, e2 + e4. Then we substitute that in:

e1 = ?x + ?y
e3 = ?x - ?y
e2 + e4 = e2 + e4

That last constraint can now be discarded. So now we have two constraints to solve for ?x and ?y. Well, here the linear solver comes in. First, we have the formula ∃ ?y, e1 = ?x + ?y /\ e3 = ?x - ?y. This can quantifier-eliminate into ?x = ½ e1 + ½ e3. Substituting that in, we have:

e1 = ½ e1 + ½ e3 + ?y
e3 = ½ e1 + ½ e3 - ?y

You do a bit of rearrangement and you get ?y = ½ e1 - ½ e3. That gives us our final binding:

?a = e1 /\ ?b = e2 /\ ?c = e3 /\ ?d = e4 /\
?e = e2 + e4 /\ ?x = ½ e1 + ½ e3 /\ ?y = ½ e1 - ½ e3

So that's a successful pattern match; note that all of these variables are of the appropriate type to be used in later terms.

Note here that to solve for ?x I couldn't work one constraint at a time. I'm not sure why that is. So I suppose that either you need theories to support quantifier elimination from conjunctions of patterns, or you need to do something else smart.

One thing I want to draw attention to, though. Here I was careful to introduce only as many linear variables as I had linear patterns. If, instead of ?x + ?y and ?x - ?y, I had ?x + ?y and ?x - ?z, the solution would have had a free variable that I was unable to finding a binding for. This would be bad, in a different way from "no pattern match". Perhaps in this case you'd want to switch to Phil Zucker style "bottom-up e-matching", where you try every single e-class? Or maybe you want to warn the user about a bad pattern? Or maybe you just sub in a "default" value like 0. In my mind this is still an open question; perhaps it depends on whether that free variable appears on the right-hand side.

By the way, if you do end up supporting free variables somehow, then I think we can do the pattern matching one constraint at a time; you end up not finding any constraints on ?x by considering e1 = ?x + ?y and e3 = ?x - ?y separately, but then you find both ?y = ?x - e1 and ?y = e3 - ?x, and conjoining those two gives you ?x - e1 = e3 - ?x, which you can then solve to "fill in" the variable you skipped. Not sure if this is a good way to work in general or not.

I'd love to see this implemented; I'm not sure it's useful but it should be pretty fast, since actually solving linear equations should be very, very fast.

By the way, one thing I think egraphs modulo theory give you, which is very nice, is the ability to use a theory more than once. For example, you might use two copies of the linear real theory to handle both sums and products, though with products you do need to be careful with zeros; I'm not sure what exactly to do with them.

Here's a weirder thing you could do. The Omega test is a (partial) quantifier elimination procedure for linear integer arithmetic, and one nice feature it has is that eliminating existential variables from conjunctions of linear equations yields more conjunctions of linear equations, plus divisibility criteria. I'm not sure what to do with divisibility criteria in general. But here's one crazy use case.

Suppose you're using the techniques described here to think about products; that is, a "linear combination" is a monomial like 3 x^2 y^3. One issue you might be worried about is the fact that you now might have to consider terms like x^(1/2), and those terms might be meaningless if, later on, you union x with -1. But terms like x^(1/3) are fine. So here's what you do: you use the Omega test, which is totally bog-standard substitution for linear equalities, except it keeps track of divisibility on the side. Then, if you ever find yourself looking at a constraint like 3 | ?x, that's OK: every value has a cube root. But if you find a constraint like 2 | ?x, that's bad: not every value has a square root. But maybe this can be rewritten to another constraint like ?x = fabs(?y), which is then injected back into the query and solved normally.

Ok, in that last bit I'm talking about something a bit crazy, where instead of constraints just substituting and eliminating quantifiers, you're also spawning new constraints to solve. Fine. Hopefully the more sober treatment of LIN-graphs was more to your liking then. Until next time!