Challenge problem: improving proof performance via smart query compilation

[oflatt]

A tough benchmark that maybe @yihozhang can tackle or others!

I’ve attached two files from my new PR: #850. One is math-micro-term.egg and one is math-micro-proof.egg, math-microbenchmark with the term encoding vs full proof encoding. math-micro-term.egg is very quick for the rules the user wrote, comparable to normal egglog:

Ruleset : search 0.537s, merge 0.318s, rebuild 0.644s

On the other hand, adding proof lookups makes math-micro-proof.egg very slow for user rules:

Ruleset : search 30.739s, merge 1.004s, rebuild 3.386s

math-micro-proofs.txt

math-micro-terms.txt

You can ignore the timing for the rebuilding and unionfind rulesets for now, the main challenge is the user rules. Would be interested if you give an analysis of how we might speed up those though (I do have some ideas).

[coderabbitai]

🔗 Related PRs

#828 - Rust rule benchmarking [merged]
#837 - Rip out backend proofs (unused) in favor of the new proof encoding [closed]
#848 - Rebuild only when union find changes [merged]

---

📝 Issue Planner

_{Check the box below or use the @coderabbitai plan command to generate an implementation plan and prompts that you can use with your favorite coding assistant.}

• Create Plan

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!

[ezrosent]

To recap and flesh out what we discussed on Friday,

The Problem

Proofs transform a number of queries like:

(rule ((= x (Add a b)) (= y (Add c x))) ( ... body ...))

To ones like

(rule ((= x (Add a b)) (= y (Add c x))
          (= prf1 (AddProof x))
          (= prf2 (AddProof y))
         ) ( ... body doing things with proofs ...))

Doing so adds two new conjuncts to the LHS of the rule, meaning that compared to the original rule we:

• Run 2 extra queries for semi-naive evaluation
• Need to construct, scan, build new indices on AddProof.

Doing so adds a good amount of overhead that is not necessary. In the case of proofs, we should be able to execute the new rule as follows:

(rule ((= x (Add a b)) (= y (Add c x)))
         ((let prf1 (AddProof x))
          (let prf2 (AddProof y))
           ... rest of body ...))

This fixes both problems. It keeps the same number of conjuncts as before, and it also does a raw lookup on AddProof, which avoids building or scanning an extra index (egglog tables support lookups of this form natively without any auxiliary data-structures).

This kind of let in the RHS of the rule is supported in the backend, but it is not safe to perform in general as lets on the RHS of the rule effectively opt the function being read out of semi-naive tracking. The result is a rule that could potentially miss updates. We could build a custom desugaring just for proofs that relies on invariants about proofs always being populated in lock-step, but that would narrow the scope of this optimization and also be error-prone.

Sketches of a Fix

Instead we discussed adding a few changes to the query planner to natively and safely support this pattern.

Support raw lookups without building extra indexes

Variables like prf1 and prf2 have a very specific form where they are (a) functional dependencies of other variables and (b) not relied on by more than one other atom in the rule. Yihong's new query planner in #785 already detects the related notion of "ear" in a query plan and we could detect such variables by extending that code. Instead of these variables being populated by a separate trie node, they can instead be directly bound once by a new Lookup instruction once all their dependencies have been satisfied.

This should generalize to "linear" portions of a query that look like (= prf (F (G (H x)))): these could become iterated lookups.

Flexible Seminaive

Seminaive evaluation works as follows for a rule that last ran at a timestamp t with conjuncts from A(...), B(...), and C(...). With 3 conjuncts we run 3 queries

• A[t..], B[0..t], C[0..t]
• A, B[t..], C[0..t]
• A, B, C[t..]

A different, mostly valid (under certain assumptions --- see the note below), set of queries is:

• A[t..], B[0..t], C
• A, B[t..], C

This effectively "turns off" seminaive for C but keeps it around for A and B.

Intuitively, if C is only ever targeted by Lookup operations in the query plan, then we should treat it in the latter way and not the former because we are never doing any work that depends on the size of C in the body of the query. In this way we should have a simple, static means of detecting when turning off seminaive would be profitable.

NOTE: there are some edge-cases here where if C contains FDs that changed in a way "not detected by other unions", then we'd want to re-run the whole query. Everything should still work but we'll want to have some careful tests here.

For a query like A(a), B(b), C(a, b)=x, if we did a destructive update from C(a,b) = 1 to C(a,b) = 2, the original set of 3 queries would detect this and potentially allow us to scan only new entries in C, while the new queries will fail to detect this if neither A nor B have changed. Proofs cannot have this behavior because every set also generates a union, so running this query again will be redundant.

Other thoughts

It's at least plausible that other non-proof queries could benefit from this optimization too. I'm excited to see if these optimizations will help once the Tree Decomposition PR is in.