Mind the Abstraction Gap: Bringing Equality Saturation to Real-World ML Compilers
summary
Equality saturation is integrated into XLA with a C++/Rust egraph pass that replaces brittle local rewrites with a global cost model, yielding an average 3.45% speedup and up to 56.26% on NasRNN (CPU). Extraction uses an ILP augmented by a fusion bonus so chosen nodes favor kernel fusion; large graphs are handled by segmentation into independent egraphs with tunable segment size; empirical results show practical scalability and fewer hand-tuned heuristics.
AI-written summary · about
chapters (12)
- 0:00Introduction
- 0:22What is equality saturation
- 1:48How egraphs enable equality saturation
- 3:09ML compilers and XLA
- 4:34Mapping XLA to egraphs
- 5:15Scaling with segmentation
- 6:40Extraction challenges and fusion
- 8:46Replacing local heuristics
- 9:29Empirical evaluation
- 10:32Key conclusions
- 10:55Q&A: Fusion vs cost model
- 11:38Q&A: Segmentation heuristics
related talks
30megg: Fast and Extensible Equality Saturation
26mEquality Saturation Guided by Large Language Models
20mMachine Learning Guided Equality Saturation
21mImproving Equality Saturation for EDA via Semantic E-Graphs
19meqsat: An Equality Saturation Dialect for Non-destructive Rewriting
20mEquality Saturation and Joins
11mEnabling Composition of Kernels and Compilers
20mVectorizing Sparse Coiteration for Two-finger Loop Structure (Extended Abstract)