Equality Saturation in MimIR
Equality Saturation is a compiler optimization technique that is primarily used to solve the Phase-Ordering Problem for compiler optimization passes. It utilizes E-Graphs to simultaneously represent a set of equivalent program terms according to a set of rewrite-rules and find the most optimal one according to a cost heuristic. This repository contains Equality Saturation implementations in egg and slotted-egraphs as a plugin for the functional higher-order intermediate representation MimIR.
You may use this plugin through the MimIR C++ API or its textual representation Mim. Consider the following lightweight examples to get started. The examples both perform the same optimization:
- Define a rewrite-rule
?n + 0 => ?n - Define a term
fun(x: Nat): Nat = return (x + 0); - Perform equality saturation in
slotted-egraphs - Extract an optimal term by smallest
AstSize
#include <fstream>
#include <mim/driver.h>
#include <mim/ast/parser.h>
#include <mim/pass/optimize.h>
#include <mim/util/sys.h>
#include <mim/plug/eqsat/eqsat.h>
using namespace mim;
using namespace mim::plug;
int main(int, char**) {
try {
auto driver = Driver("eqsat");
auto& w = driver.world();
driver.log().set(&std::cerr).set(Log::Level::Debug);
ast::load_plugins(w, View<std::string>{"core", "ll", "eqsat"});
// rule foo (x: Nat): %core.nat.add (x, 0) => x;
auto foo = w.mut_rule(w.type_nat())->set("foo");
auto x = foo->var()->set("x");
auto lhs = w.call(core::nat::add, w.tuple(x, lit_nat(0)))
auto rhs = x;
foo->set_lhs(lhs);
foo->set_rhs(rhs);
foo->set_guard(w.lit_tt());
// Quickly define config values
eqsat_config(
w,
eqsat::slotted,
eqsat::AstSize,
eqsat_rulesets(eqsat::standard),
eqsat_rules(foo),
);
// fun extern main(x: Nat): Nat = return %core.nat.add (x, 0);
auto main = w.mut_fun({w.type_nat()}, {w.type_nat()})->set("main");
auto x = main->var(2, 0)->set("x");
auto ret = main->var(2, 1);
main->app(false, ret, x);
main->externalize();
// Equality saturation and code gen are performed here
optimize(w);
sys::system("clang eqsat.ll -o eqsat -Wno-override-module");
std::println("exit code: {}", sys::system("./eqsat"));
} catch (const std::exception& e) {
std::println(std::cerr, "{}", e.what());
return EXIT_FAILURE;
} catch (...) {
std::println(std::cerr, "error: unknown exception");
return EXIT_FAILURE;
}
return EXIT_SUCCESS;
}plugin core;
plugin eqsat;
// You can define your own syntactic rewrite-rules here
rule foo (x: Nat): %core.nat.add (x, 0) => x;
lam extern _config() =
%eqsat.config (
// Specifies whether the plugin should use its egg or slotted-egraphs backend
%eqsat.slotted,
// Defines the cost function that should be used for term extraction
%eqsat.AstSize,
// Specifies a set of rules directly implemented in egg or slotted-egraphs
// To implement and use your own ruleset, follow the instructions under **Rulesets**.
%eqsat.rulesets (%eqsat.normalize),
// To use the rule 'foo' that we defined above for equality saturation
%eqsat.rules (foo),
// Here you may provide two terms to assert whether term A can reach term B in a number of steps
%eqsat.reaches (term_A, term_B, 10),
// Here you may select specific terms that should be rewritten
// When providing an empty tuple, no terms will be rewritten
%eqsat.select (),
);
fun extern main(x: Nat): Nat =
return %core.nat.add (x, 0);
Clone the mimir repository
git clone --recursive https://github.com/mimir/mimir.gitClone the eqsat repository
cd mimir/extra
git clone https://github.com/ashiven/eqsat.git
cd ..Ensure that Rust and Cargo are installed
curl https://sh.rustup.rs -sSf | shBuild the project
cmake -S . -B build -DBUILD_TESTING=ON -DMIM_BUILD_EXAMPLES=ON
cmake --build build -j$(nproc)You may want to define a set of rewrite-rules that are more complex than the syntactic rewrite-rules that can be defined in MimIR. In this case, you should follow the implementation guide below on adding a set of rules and a new analysis directly in egg or slotted-egraphs.
Automatically generate all of the boilerplate code required to integrate your ruleset with the eqsat plugin
python ./scripts/new_ruleset.py egg MyRulesDefine your ruleset in src/egg/rulesets/myrules.rs
use crate::egg::{Mim, analysis::AnalysisData, analysis::MimAnalysis};
use egg::{EGraph, Rewrite, Pattern, DidMerge, Id};
pub fn rules() -> Vec<Rewrite<Mim, MimAnalysis>> {
let rules = vec![
my_rule(),
];
rules
}
fn my_rule() -> Rewrite<Mim, MimAnalysis> {
let pat: Pattern<Mim> = "(app %foo.bar ?baz)".parse().unwrap();
let outpat: Pattern<Mim> = "?baz".parse().unwrap();
Rewrite::new("my-rule", pat, outpat).unwrap()
}
pub type MyRulesData = ();
pub struct MyRulesAnalysis;
impl MyRulesAnalysis {
pub fn make(_eg: &mut EGraph<Mim, MimAnalysis>, _enode: &Mim, _id: Id) -> AnalysisData {
AnalysisData::default()
}
pub fn merge(_l: &mut AnalysisData, _r: AnalysisData) -> DidMerge {
DidMerge(false, false)
}
pub fn modify(_eg: &mut EGraph<Mim, MimAnalysis>, _id: Id) {}
}To define your own cost function for term extraction, follow the steps below.
Automatically generate all of the boilerplate code required by the eqsat plugin
python ./scripts/new_cost.py egg MyCostDefine your new cost function in src/egg/cost.rs
#[derive(Debug)]
pub struct MyCost;
impl CostFunction<Mim> for MyCost {
type Cost = usize;
fn cost<C>(&mut self, enode: &Mim, mut costs: C) -> Self::Cost
where
C: FnMut(Id) -> Self::Cost,
{
enode.fold(1, |sum, id| sum.saturating_add(costs(id)))
}
}This library also exposes its methods in a C++ FFI, which was required to integrate it into the MimIR plugin system. The following documents the signatures generated for these methods via CXX along with a short description of what they do.
/**
* Rewrites an sexpr in `egg` format
*
* sexpr: a symbolic expr in `egg` format (emitted by the `mim` compiler via `--output-sexpr`)
* selected: optionally, a list of identifiers for terms that should be rewritten
* rulesets: a list of identifiers of rulesets that should be used for rewriting (see src/egg/rulesets)
* cost_fn: a cost function that should be used for term extraction (currently only AstSize and AstDepth)
*/
rust::Vec<RecExprFFI> eqsat_egg(rust::Str sexpr, OptionSelected selected, rust::Vec<RuleSet> rulesets, CostFn cost_fn);/**
* Rewrites an sexpr in `slotted-egraphs` format
*
* sexpr: a symbolic expr in `slotted-egraphs` format (emitted by the `mim` compiler via `--output-sexpr-slotted`)
* selected: optionally, a list of identifiers for terms that should be rewritten
* rulesets: a list of identifiers of rulesets that should be used for rewriting (see src/slotted/rulesets)
* cost_fn: a cost function that should be used for term extraction (currently only AstSize)
*/
rust::Vec<RecExprFFI> eqsat_slotted(rust::Str sexpr, OptionSelected selected, rust::Vec<RuleSet> rulesets, CostFn cost_fn);/**
* Uses `slotted-egraphs` to prove whether two terms are equivalent
*
* sexpr: a symbolic expr in `slotted-egraphs` format (emitted by the `mim` compiler via `--output-sexpr-slotted`)
* rulesets: a list of identifiers of rulesets that should be used for rewriting (see src/slotted/rulesets)
* start_name: an identifier for the starting term
* end_name: an identifier for the end term that should be reached via equality saturation
* max_steps: the maximum number of iterations in which the start term should reach the end term
*/
bool reaches_egg(rust::Str sexpr, rust::Vec<RuleSet> rulesets, rust::Str start_name, rust::Str end_name, std::size_t max_steps);/**
* Uses `egg` to prove whether two terms are equivalent
*
* sexpr: a symbolic expr in `egg` format (emitted by the `mim` compiler via `--output-sexpr`)
* rulesets: a list of identifiers of rulesets that should be used for rewriting (see src/slotted/rulesets)
* start_name: an identifier for the starting term
* end_name: an identifier for the end term that should be reached via equality saturation
* max_steps: the maximum number of iterations in which the start term should reach the end term
*/
bool reaches_slotted(rust::Str sexpr, rust::Vec<::RuleSet> rulesets, rust::Str start_name, rust::Str end_name, std::size_t max_steps);/**
* Pretty-prints an sexpr in `egg` format
*
* sexpr: a symbolic expr in `egg` format (emitted by the `mim` compiler via `--output-sexpr`)
* line_len: the maximal line length after which the sexpr continues on a new line
*/
rust::String pretty_egg(rust::Str sexpr, std::size_t line_len);/**
* Pretty-prints an sexpr in `slotted-egraphs` format
*
* sexpr: a symbolic expr in `slotted-egraphs` format (emitted by the `mim` compiler via `--output-sexpr-slotted`)
* line_len: the maximal line length after which the sexpr continues on a new line
*/
rust::String pretty_slotted(rust::Str sexpr, std::size_t line_len);/**
* Pretty-prints an sexpr represented by a Vec<RecExprFFI>
*
* sexprs: a vector of symbolic expressions in RecExprFFI format (the result of equality saturation)
* line_len: the maximal line length after which the sexpr continues on a new line
*/
rust::String pretty_ffi(rust::Vec<RecExprFFI> sexprs, std::size_t line_len);Please feel free to submit a pull request or open an issue.
- Fork the repository
- Create a new branch:
git checkout -b feature-name. - Make your changes
- Push your branch:
git push origin feature-name. - Submit a PR
This project is licensed under the MIT License.