EGraph for RISC-V Assembly: Making Optimization More Declarative

Compilers and assemblers usually optimize through hand-written passes: pattern checks, ad hoc rewrites, and a lot of careful bookkeeping to avoid breaking correctness. In this project, I took a different route. I built an e-graph-based optimization module for RISC-V assembly in src/egraph/rvasm.rs, so assembly optimization can be expressed as rules instead of scattered special cases.

At a high level, the module treats assembly as a language that can live inside an equality saturation engine. That means equivalent forms are collected into the same equivalence classes, rewrites are applied until saturation, and the best version is chosen later by cost. The result is a more declarative pipeline: instead of encoding one fixed transformation path, the engine explores many possible forms and lets the extractor pick the best one.

The core design is simple but flexible. The assembly language is represented by two node types: individual instructions and whole program blocks. Instructions carry the original parsed assembly item, while program blocks hold sequences of e-class ids. That structure lets the engine reason about both local instruction-level equivalences and larger block-level composition.

The instruction node also defines the behavior that matters for optimization:

• Structural matching, so equivalent instructions can be recognized.
• A small cost model, so cheaper forms like add-immediate or add can be preferred appropriately.
• Display logic, so extracted results can be rendered back into readable assembly.

This is the foundation that makes the rest of the system work. Once assembly is expressed as a language inside the e-graph, it becomes possible to compare alternative forms without committing to a single rewrite path too early.

The main advantage of using an e-graph is that it keeps all equivalent variants available at once. For assembly, that matters because many optimizations are not strictly linear. A transformation like strength reduction can enable another simplification, and a later rewrite may expose an even better form. Traditional pass pipelines can miss those opportunities depending on pass ordering.

With the e-graph approach, those relationships are represented naturally. The module can explore transformations such as:

• rewriting multiplication by powers of two into shifts,
• simplifying instructions with neutral elements,
• handling conditional rewrites when analysis proves a property,
• generating dynamic rewrites based on matched values.

The module’s tests and demos show this style clearly in src/egraph/tests.rs and src/tests/test_egraph.rs, where the behavior is pinned down with concrete cases like constant folding, congruence merging, and rewrite scheduling.