Equality Saturation in MimIR
Equality Saturation is a compiler optimization technique that uses E-Graphs to search large transformation spaces. This repository contains an Equality Saturation plugin for the MimIR compiler framework.
Warning
MimIR 0.3 introduced several breaking changes that broke parts of the plugin that previously worked. (See unsupported test cases in the lit directory)
You can use this plugin through the MimIR C++ API or its textual representation Mim. Consider the following lightweight examples to get started. The examples perform the same transformation:
- 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 <fe/sys.h>
#include <mim/driver.h>
#include <mim/ast/parser.h>
#include <mim/phase/optimize.h>
#include <mim/plug/eqsat/eqsat.h>
using namespace mim;
using namespace mim::plug;
int main(int, char**) {
auto driver = Driver("eqsat");
try {
auto& w = driver.world();
driver.log().set(&std::cerr).set(Log::Level::Debug);
ast::load_plugins(w, View<std::string>{"core", "ll", "eqsat", "sexpr"});
// 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);
fe::sys::system("clang eqsat.ll -o eqsat -Wno-override-module");
std::println("exit code: {}", fe::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 and sexpr repositories
cd mimir/extra
git clone https://github.com/ashiven/eqsat.git
git clone https://github.com/ashiven/sexpr.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)To define a set of rewrite-rules directly in egg or slotted-egraphs, follow the implementation guide below.
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.
Generate all of the boilerplate code required by the eqsat plugin
python ./scripts/new_cost.py egg MyCostDefine your 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.