Lab 17 ★ · A small equality-saturation engine¶

Chapter: 17 · SSA-Based Scalar Optimizations · Lessons: 17.8 (Definitions 17.8.2–17.8.3, Algorithm 17.8.4), e-graphs as an IR in Lesson 8.5 · Time: 6–10 hours · Tests: ./course test 17 (label ch17, suite ch17.lab: egraph.test) · Optional (★)

Goal¶

Build an e-graph and run equality saturation on it, in rounds, the way egg does (Algorithm 17.8.4). The e-graph uses union-find over e-class ids, a hashcons of canonical e-nodes and deferred rebuilding. Each round has a read phase that e-matches every rule, a write phase that instantiates and merges, and a rebuild that restores congruence. At the end you extract a smallest term. Then use the engine to see Lesson 17.8's claims happen: - the running example's facts, written as rules, make t − u + x collapse to x/1 in 4 rounds; - ring identities saturate on some terms and grow without bound on others; - two terms can be proven equal by ending up in one class.

Your class, e-node and round counts must match egg 0.11's on the lab inputs, so the round structure is fixed (R2). How you store the e-graph is up to you.

Requirements¶

Terms, patterns, rules (provided parser, include/lab17eg/Sexp.h): - term ::= atom | "(" op term+ ")", where an atom is a symbol or an integer, compared as text. - Patterns may also contain variables ?name. - A rules file has one rule name: lhs => rhs per line, with # comments. The left-hand side is not a variable, and every variable of the right-hand side occurs on the left.

G1 (e-graph). Implement Definition 17.8.2:
e-class ids with union-find;
e-nodes op(c1, …, ck) whose children are class ids;
a hashcons from canonical e-nodes (children replaced by find) to classes.

Adding a term adds its subterms bottom-up and reuses existing e-nodes (hashconsing). After every rebuild, the congruence invariant holds: no two classes contain the same canonical e-node. - G2 (rounds). Algorithm 17.8.4, exactly: 1. Read phase. Collect every match \((\text{rule}, c, \sigma)\) of every rule's left-hand side in every class \(c\) of the e-graph as it is at the start of the round. E-matching (Definition 17.8.3) binds each variable to a class. A repeated variable must bind to the same class. 2. Write phase. For each match, add \(r\sigma\) and merge its class with \(c\). 3. Rebuild. Re-canonicalize and merge congruent e-nodes until the invariant holds.

A round changes the e-graph if it created an e-node or merged two different classes. After the rebuild, stop with: - saturated if the round changed nothing; - otherwise node-limit if the number of e-nodes exceeds Limits::MaxNodes; - otherwise iteration-limit if this was round Limits::MaxIterations.

Iterations counts the rounds run, including the last one that changed nothing. - G3 (counts). Classes is the number of e-classes at the end. Nodes is the number of distinct canonical e-nodes at the end, i.e. the hashcons size after the last rebuild. - G4 (extraction). Best is a term of least size (1 per operator or leaf) among the terms the root's class represents. Cost is its size. Compute it bottom-up to a fixed point (Bellman–Ford on classes, Algorithm 17.8.4 Extract): a class's cost is the minimum over its e-nodes of 1 + the children's costs. Ties may be broken in any way. The tests accept every least-size term where there is more than one. - G5 (equality). provablyEqual(A, B, …) adds both terms to one e-graph, runs G2 under the same limits, and returns whether find(A) == find(B) at the end. - Errors. saturate and provablyEqual return an error for a rule whose left-hand side is a variable. The provided parser already rejects that, so this matters only for rules built in code. - Performance. Every lab input finishes in well under a second. The largest, ring.rules on (* (+ a b) (+ c d)) with --iters 4, has 46 e-nodes. E-matching by backtracking over class members is fine at this size.

The contract¶

// labs/ch17-egraph/include/lab17eg/EGraph.h  (provided; do not change)
namespace lab17eg {
struct Limits { unsigned MaxIterations = 30; unsigned MaxNodes = 10000; };
struct Result {
  sexp::Term Best; unsigned Cost = 0; unsigned Classes = 0; unsigned Nodes = 0;
  unsigned Iterations = 0; std::string Stop;   // "saturated" | "iteration-limit" | "node-limit"
};
std::expected<Result, std::string> saturate(const sexp::Term &T, const std::vector<sexp::Rule> &Rules,
                                            const Limits &L = {});                 // G1-G4
std::expected<bool, std::string> provablyEqual(const sexp::Term &A, const sexp::Term &B,
                                               const std::vector<sexp::Rule> &Rules,
                                               const Limits &L = {});              // G5
}

Input and output formats¶

usage: ch17-egraph --rules <file> [--iters N] [--nodes N] (<term> | --equal <term> <term>)

For a term, it prints six lines, and for --equal it prints equal or not proven:

$ ch17-egraph --rules labs/ch17-egraph/inputs/div.rules '(/ (* a 2) 2)'
best: a
cost: 1
classes: 4
nodes: 10
iterations: 4
stop: saturated

A rules file (inputs/running.rules):

sccp-x: x => 1
gvn-j: j => i
sub-self: (- ?a ?a) => 0
add-zero-l: (+ 0 ?a) => ?a
comm-add: (+ ?a ?b) => (+ ?b ?a)

Provided infrastructure¶

File	What it gives you
`include/lab17eg/EGraph.h`	the contract
`include/lab17eg/Sexp.h`, `provided/Sexp.cpp`	terms, patterns and the rules-file reader (`parseTerm`, `parseRules`, `Term::size`, `Term::str`)
`tools/EGraphMain.cpp`	`ch17-egraph`
`inputs/div.rules`	the egg paper's and Lesson 8.5's example (`div-canc` and `div-self` hold for mathematical integers with a nonzero divisor, not for machine arithmetic)
`inputs/ring.rules`	identities of the integers modulo \(2^{64}\): commutativity, associativity, neutral and absorbing elements, distributivity both ways, `* 2` → `<< 1`
`inputs/running.rules`	the running example's SCCP and GVN facts as rules (Lesson 17.8 §3)
`src/Stub.cpp`	the two contract functions, stopping with `TODO(ch17)`

What the tests check¶

All in tests/ch17/lab/egraph.test (suite ch17.lab). The counts equal egg 0.11's (Runner with SimpleScheduler, AstSize extraction) on the same rules and terms. The only exception is the node-limit case, because egg checks its node limit inside a round.

Case	Checks
`div.rules`, `(/ (* a 2) 2)`	best `a`, cost 1, 4 classes, 10 nodes, 4 rounds, saturated (G1–G4)
`running.rules`, `(+ (- (* i 4) (* j 4)) x)`	best `1` or `x`, cost 1, 5 classes, 10 nodes, 4 rounds, saturated
`ring.rules`, `(+ (* a 2) (* a 3))`	best `(* a (+ 2 3))` or `(* a (+ 3 2))`, cost 5, 8 classes, 15 nodes, 3 rounds, saturated
`ring.rules --iters 4`, `(* (+ a b) (+ c d))`	cost 7, 19 classes, 46 nodes, 4 rounds, `iteration-limit`
`ring.rules --iters 3`, `(- (* (+ x 1) 2) (* 2 (+ 1 x)))`	best `0`, cost 1, 10 classes, 26 nodes
`ring.rules --nodes 25`, `(* (+ a b) (+ c d))`	15 classes, 30 nodes, 3 rounds, `node-limit` (the limit is checked after the round)
`--equal` (G5)	`(* (+ a b) c)` and `(+ (* c b) (* a c))` are `equal` within 8 rounds; `(* a b)` and `(+ a b)` are `not proven`
soundness	with `ring.rules --iters 5`, the best term for `(+ (* (+ a b) (+ a b)) (- (* a 0) (* b 2)))` evaluates to the same value as the input on 200 random 64-bit assignments (`Inputs/eval_terms.py`)

Milestones¶

G1 with no rules: ch17-egraph --rules /dev/null '(+ (* a 2) (* a 2))' gives 4 classes and 4 nodes (hashconsing shares the two products), 1 round, saturated.
G2 with e-matching but a naive rebuild (re-canonicalize everything until nothing changes): div.rules gives 4 / 10 / 4.
G3–G4: all counts in the table. Then run ctest --test-dir build/<preset> -R ch17.lab.
G5 and the limits.

Hints¶

Hint 1 — where to start

Lesson 17.8 §3 lists the running example's e-graph after every round, and Lesson 8.5 §3 does the same for div.rules. Reproduce them with a debug print of the classes before you worry about the limits. Represent an e-node as (op string, vector of class ids) with operator<=>, so that a std::map can be the hashcons.

Hint 2 — the key idea

The read and write phases must be separated. Collect all matches first, then apply them. Applying a match immediately changes what later rules see in the same round, and the counts then differ from egg's.
Rebuilding is what finds new equalities by congruence. After j ≡ i, (* j 4) and (* i 4) have the same canonical form and must be merged, which can make their parents congruent in turn. Repeat until nothing merges.
Extraction needs a fixed point, not one pass, because classes can be cyclic (e6 contains (+ e5 e6) in the running example).

Hint 3 — a design sketch

Use a std::vector<unsigned> parent array for union-find, and keep the smaller id as the root so the output is deterministic. Keep a std::map<ENode, unsigned> hashcons, and rebuild it from scratch in rebuild() until no key collides.
Build std::map<unsigned, std::vector<ENode>> classes after each rebuild, for e-matching and extraction. E-matching is a recursive function (pattern, class, partial substitution) that returns every consistent substitution.
Set a Changed flag in add (on a new e-node) and in merge (on two different roots).

Bugs the tests catch most often: - counting e-nodes before the rebuild, which gives too many because duplicates are not yet merged; - stopping one round early (a round that changes nothing still counts); - matching a repeated variable, as in (- ?a ?a), against two different classes.

Stretch goals ★¶

Parent lists and egg's incremental rebuild [WNW+21, §3]: re-canonicalize only the parents of merged classes. Count the e-nodes visited per rebuild before and after.
An e-class analysis for constants (egg §4): attach optional<int64_t> to each class, merge on union, and add the constant as an e-node when it becomes known. Then (+ 2 3) is folded without a rule.
A backoff scheduler: ban a rule for \(2^k\) rounds after it matched more than a threshold. Compare what (* (+ a b) (+ c d)) reaches within the same node budget.