eqsat: An Equality Saturation Dialect for Non-destructive Rewriting

2025-05-14

AI Summary

Traditional equality saturation suffers from high transformation overhead and difficulty in maintaining persistent e-graph state due to the separation between e-graph libraries and IRs in compiler infrastructures. This paper introduces *eqsat*, a native equality saturation dialect embedded directly into MLIR. Its core innovation is modeling the e-graph as an intrinsic MLIR intermediate representation, enabling non-destructive, composable updates to e-graph state throughout the entire compilation pipeline. By deep integration at the IR level, eqsat allows constructive passes and e-graphs to co-evolve, while reusing standard MLIR optimizations—such as common subexpression elimination (CSE)—to accelerate e-matching and extraction, thereby eliminating cross-system data conversion. This design enables any domain-specific IR to seamlessly integrate into the standard MLIR infrastructure for equality saturation optimization. Empirical results demonstrate significant improvements in expressiveness and end-to-end efficiency for hardware synthesis and program synthesis tasks.

Technology Category

• Constraint Satisfaction and Optimization: Satisfiability Modulo Theories
• Machine Learning: Hardware-aware ML
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Proposes eqsat dialect for native e-graph representation in MLIR
• Enables equality saturation in compiler IR without external libraries
• Facilitates constructive compiler passes maintaining e-graph state throughout compilation

Innovation

• Native e-graph representation in MLIR
• Integrates equality saturation in compiler IR
• Avoids translation overhead with MLIR passes