SmoothE: Differentiable E-Graph Extraction

E-graphs have gained increasing popularity in compiler optimization, program synthesis, and theorem proving tasks. They enable compact representation of many equivalent expressions and facilitate transformations via rewrite rules without phase ordering limitations. A major bene!t of using e-graphs is the ability to explore a large space of equivalent expressions, allowing the extraction of an expression that best meets certain optimization objectives (or cost models). However, current e-graph extraction methods often face unfavorable scalability-quality trade-o"s and only support simple linear cost functions, limiting their applicability to more realistic optimization problems.

In this work, we propose SmoothE, a di"erentiable e-graph extraction algorithm designed to handle complex cost models and optimized for GPU acceleration. More speci!cally, we approach the e-graph extraction problem from a probabilistic perspective, where the original discrete optimization is relaxed to a continuous di"erentiable form. This formulation supports any di"erentiable cost functions and enables e#cient searching for solutions using gradient descent. We implement SmoothE in PyTorch to leverage the advancements of the modern machine learning ecosystem. Additionally, we introduce performance optimization techniques to exploit sparsity and data parallelism. We evaluate SmoothE on a variety of realistic e-graphs from !ve di"erent applications using three distinct cost models, including both linear and non-linear ones. Our experiments demonstrate that SmoothE consistently achieves a favorable trade-o" between scalability and solution quality.