eqsat: An Equality Saturation Dialect for Non-destructive Rewriting

TL;DR

The paper presents eqsat, a dialect that embeds $e$-graphs and equality-saturation machinery directly into an MLIR-based compiler IR. By introducing ops for $e$-classes, $e$-graphs, and a region-bound eqsat.egraph, it enables non-destructive, stateful rewriting across arbitrary domain-specific IRs while leveraging existing MLIR passes (e.g., CSE) and pattern infrastructure (PDL). The approach facilitates handling of cycles via graph regions and supports cross-cutting rewrites through structured control flow, with a prototypical open-source implementation in xDSL. This integration reduces translation overhead from external libraries and unlocks new rewriting opportunities across dialects and abstraction levels, with future work on cost models and region-aware pattern matching.

Abstract

With recent algorithmic improvements and easy-to-use libraries, equality saturation is being picked up for hardware design, program synthesis, theorem proving, program optimization, and more. Existing work on using equality saturation for program optimization makes use of external equality saturation libraries such as egg, typically generating a single optimized expression. In the context of a compiler, such an approach uses equality saturation to replace a small number of passes. In this work, we propose an alternative approach that represents equality saturation natively in the compiler's intermediate representation, facilitating the application of constructive compiler passes that maintain the e-graph state throughout the compilation flow. We take LLVM's MLIR framework and propose a new MLIR dialect named eqsat that represents e-graphs in MLIR code. This not only provides opportunities to rethink e-matching and extraction techniques by orchestrating existing MLIR passes, such as common subexpression elimination, but also avoids translation overhead between the chosen e-graph library and MLIR. Our eqsat intermediate representation (IR) allows programmers to apply equality saturation on arbitrary domain-specific IRs using the same flow as other compiler transformations in MLIR.

Figures (1)

• Figure 1: (left) In MLIR's default, destructive rewriting, pdl_interp. get_result and pdl_interp. get_defining_op are each other's inverse. (right) In equality saturation, this is not the case because each value comes from one of multiple equivalent operations.