Equality Saturation in MimIR
Equality Saturation is a compiler optimization technique that facilitates the discovery of complex optimization sequences. It utilizes E-Graphs to simultaneously represent a group of equivalent program terms and discover the best term according to a cost heuristic. This repository contains Equality Saturation implementations in egg and slotted-egraphs as a plugin for the compiler 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>usingnamespacemim;usingnamespacemim::plug;intmain(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 valueseqsat_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 hereoptimize(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());
returnEXIT_FAILURE;
} catch (...) {
std::println(std::cerr, "error: unknown exception");
returnEXIT_FAILURE;
}
returnEXIT_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
usecrate::egg::{Mim, analysis::AnalysisData, analysis::MimAnalysis};use egg::{EGraph,Rewrite,Pattern,DidMerge,Id};pubfnrules() -> Vec<Rewrite<Mim,MimAnalysis>>{let rules = vec![
my_rule(),];
rules
}fnmy_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()}pubtypeMyRulesData = ();pubstructMyRulesAnalysis;implMyRulesAnalysis{pubfnmake(_eg:&mutEGraph<Mim,MimAnalysis>,_enode:&Mim,_id:Id) -> AnalysisData{AnalysisData::default()}pubfnmerge(_l:&mutAnalysisData,_r:AnalysisData) -> DidMerge{DidMerge(false,false)}pubfnmodify(_eg:&mutEGraph<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)]pubstructMyCost;implCostFunction<Mim>forMyCost{typeCost = usize;fncost<C>(&mutself,enode:&Mim,mutcosts:C) -> Self::CostwhereC: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*/boolreaches_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*/boolreaches_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.