01 - Basics of Equality Saturation#
This tutorial is translated from egglog.
In this tutorial, we will build an optimizer for a subset of linear algebra using egglog. We will start by optimizing simple integer arithmetic expressions. Our initial DSL supports constants, variables, addition, and multiplication.
# mypy: disable-error-code="empty-body" from __future__ import annotations from typing import TypeAlias from collections.abc import Iterable from egglog import * class Num ( Expr ): def __init__ ( self , value : i64Like ) -> None : ... @classmethod def var ( cls , name : StringLike ) -> Num : ... def __add__ ( self , other : NumLike ) -> Num : ... def __mul__ ( self , other : NumLike ) -> Num : ... # Support inverse operations for convenience # they will be translated to non-reversed ones def __radd__ ( self , other : NumLike ) -> Num : ... def __rmul__ ( self , other : NumLike ) -> Num : ... NumLike : TypeAlias = Num | StringLike | i64Like
The signature here takes NumLike not Num so that you can write Num(1) + 2 instead of Num(1) + Num(2) . This is helpful for ease of use and also for compatibility when you are trying to create expressions that act like Python objects which perform upcasting.
To support this, you must define conversions between primitive types and your expression types. When a value is passed into a function, it will find the type it should be converted to and transitively apply the conversions you have defined:
converter ( i64 , Num , Num ) converter ( String , Num , Num . var )
Now, let’s define some simple expressions.
egraph = EGraph () x = Num . var ( "x" ) expr1 = egraph . let ( "expr1" , 2 * ( x * 3 )) expr2 = egraph . let ( "expr2" , 6 * x )
You should see an e-graph with two expressions.
... continue reading