E-Graph Explained

Table of Contents

E-graph is a compact data structure that maintains a set of term rewriting rules. One can specify a group of rewriting rules, and use e-graph to search how to rewrite an expression to a desired form w/ some good property (e.g. minimal ops, simplest) by designing a cost model on "e-nodes".

Why E-graph

We must notice that the rewriting rules are not commutative, which means the rewrite order matters!

In Machine Learning compilers, a common optimization is kernel-fusion: fusing several small operators into a single large operators, thus reducing the kernel launching overhead and memory footprint. Such fusion can we viewed as term-writing on computational graphs, a simple fusion strategy is defining a cost model on operators, and always fuse operators towards reducing the overall cost. However, in this way we cannot find the global optimal computational graph, because the rewrite order matters, and the best rewrite trace might look like:

traces rewrite1 rewrite2 rewrite3 rewrite4 rewrite5
cost 1 0.9 1.2 1.1 0.5

note that rewrite3 harms total cost temporarily but lay pavement for better possible rewrites (e.g. rewrite5).

A simple idea to perform such equivalent-form search is back-tracking 1 , which means we allow some increase on cost during the search.

TODO What's E-graph

egg

Footnotes:

Author: expye(Zihao Ye)

Email: expye@outlook.com

Date: 2021-11-01 Mon 00:00

Last modified: 2022-09-23 Fri 00:04

Licensed under CC BY-NC 4.0