eqsat: An Equality Saturation Dialect for Non-destructive Rewriting

TL;DR

This paper introduces eqsat, a new MLIR dialect that natively integrates equality saturation into compiler workflows, enabling more efficient and flexible program optimization without external libraries.

Contribution

It presents eqsat, a novel MLIR dialect that embeds equality saturation directly into compiler IR, facilitating constructive passes and domain-specific optimizations.

Findings

Enables equality saturation within MLIR without external libraries.

Supports constructive compiler passes maintaining e-graph state.

Reduces translation overhead and improves optimization flexibility.

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.