EGraphs and Equality Saturation
An EGraph is an efficient data structure for representing congruence relations. EGraphs are data structures originating from theorem provers. Several projects have very recently repurposed EGraphs to implement state-of-the-art, rewrite-driven compiler optimizations and program synthesizers using a technique known as equality saturation. Metatheory.jl provides a general purpose, customizable implementation of EGraphs and equality saturation, inspired from the egg library for Rust. You can read more about the design of the EGraph data structure and equality saturation algorithm in the egg paper.
What can I do with EGraphs in Metatheory.jl?
Most importantly, the EGraph backend for Metatheory.jl allows you to create an EGraph from a starting expression, to add more expressions to the EGraph with addexpr!
, and then to effectively fill the EGraph with all possible equivalent expressions resulting from applying rewrite rules from a theory, by using the saturate!
function. You can then easily extract expressions with a cost function and an ExtractionAnalysis
.
A killer feature of egg and Metatheory.jl are EGraph Analyses. They allow you to annotate expressions and equivalence classes in an EGraph with values from a semilattice domain, and then to:
- Extract expressions from an EGraph basing from analysis data.
- Have conditional rules that are executed if some criteria is met on analysis data
- Have dynamic rules that compute the right hand side based on analysis data.
Theories and Algebraic Structures
The e-graphs backend can directly handle associativity, commutativity and distributivity, rules that are otherwise known of causing loops in symbolic computations.
comm_monoid = @theory begin
a * b => b * a
a * 1 => a
a * (b * c) => (a * b) * c
end
The Metatheory Library
The Metatheory.Library
module contains utility functions and macros for creating rules and theories from commonly used algebraic structures and properties.
using Metatheory.Library
comm_monoid = commutative_monoid(:(*), 1)
# alternatively
comm_monoid = @commutative_monoid (*) 1
Equality Saturation
We can programmatically build and saturate an EGraph. The function saturate!
takes an EGraph
and a theory, and executes equality saturation. Returns a report of the equality saturation process. saturate!
is configurable, customizable parameters include a timeout
on the number of iterations, a eclasslimit
on the number of e-classes in the EGraph, a stopwhen
functions that stops saturation when it evaluates to true.
G = EGraph(:((a * b) * (1 * (b + c))));
report = saturate!(G, t);
# access the saturated EGraph
report.egraph
# show some fancy stats
println(report);
With the EGraph equality saturation backend, Metatheory.jl can prove simple equalities very efficiently. The @areequal
macro takes a theory and some expressions and returns true iff the expressions are equal according to the theory. The following example returns true.
@areequal t (x+y)*(a+b) ((a*(x+y))+b*(x+y)) ((x*(a+b))+y*(a+b))
Configurable Parameters
EGraphs.saturate!
can accept an additional parameter of type EGraphs.SaturationParams
to configure the equality saturation algorithm. The documentation for the configurable parameters is available in the EGraphs.SaturationParams
API docstring.
# create the saturation params
params = SaturationParams(timeout=10, eclasslimit=4000)
saturate!(egraph, theory, params)
API Docs
Metatheory.EGraphs.ClassMem
— TypeAbstract type representing an EGraph
analysis, attaching values from a join semi-lattice domain to an EGraph
Metatheory.EGraphs.EGraph
— Typemutable struct EGraph
A concrete type representing an [EGraph
]. See the egg paper for implementation details
Fields
uf::DataStructures.IntDisjointSets
stores the equality relations over e-class ids
classes::Dict{Int64,EClass}
map from eclass id to eclasses
memo::Dict{ENode,Int64}
dirty::Array{Int64,1}
worklist for ammortized upwards merging
pruned::Array{Int64,1}
root::Int64
analyses::Set{Type{var"#s309"} where var"#s309"<:AbstractAnalysis}
A vector of analyses associated to the EGraph
symcache::Dict{Any,Array{Int64,1}}
a cache mapping function symbols to e-classes that contain e-nodes with that function symbol.
numclasses::Int64
numnodes::Int64
Base.merge!
— MethodGiven an EGraph
and two e-class ids, set the two e-classes as equal.
Signatures
merge!(g::EGraph, a::Int64, b::Int64) -> Int64
Methods
merge!(g, a, b)
Metatheory.EGraphs.add!
— MethodMetatheory.EGraphs.addexpr!
— MethodMetatheory.EGraphs.find
— MethodReturns the canonical e-class id for a given e-class.
Signatures
find(g::EGraph, a::Int64) -> Int64
Methods
find(g, a)
Metatheory.EGraphs.reachable
— MethodRecursive function that traverses an EGraph
and returns a vector of all reachable e-classes from a given e-class id.
Signatures
reachable(g::EGraph, id::Int64) -> Array{Int64,1}
Methods
reachable(g, id)
Metatheory.EGraphs.rebuild!
— MethodMetatheory.EGraphs.eqsat_step!
— MethodCore algorithm of the library: the equality saturation step.
Signatures
eqsat_step!(g::EGraph, theory::Array{var"#s308",1} where var"#s308"<:Rule, mod::Module, scheduler::Metatheory.EGraphs.Schedulers.AbstractScheduler, match_hist::Array{Tuple{Rule,Union{Nothing, Pattern},Base.ImmutableDict{Any,Tuple{EClass,Any}},Int64},1}, params::SaturationParams) -> Tuple{Metatheory.EGraphs.Report,EGraph}
Methods
eqsat_step!(g, theory, mod, scheduler, match_hist, params)
Metatheory.EGraphs.saturate!
— MethodGiven an EGraph
and a collection of rewrite rules, execute the equality saturation algorithm.
Signatures
saturate!(g::EGraph, theory::Array{var"#s304",1} where var"#s304"<:Rule; mod) -> Metatheory.EGraphs.Report
Methods
saturate!(g, theory; mod)