Equality Saturation for Tensor Graph Superoptimization
Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel pith:AC74XNPKrecord.jsonopen to challenge →
read the original abstract
One of the major optimizations employed in deep learning frameworks is graph rewriting. Production frameworks rely on heuristics to decide if rewrite rules should be applied and in which order. Prior research has shown that one can discover more optimal tensor computation graphs if we search for a better sequence of substitutions instead of relying on heuristics. However, we observe that existing approaches for tensor graph superoptimization both in production and research frameworks apply substitutions in a sequential manner. Such sequential search methods are sensitive to the order in which the substitutions are applied and often only explore a small fragment of the exponential space of equivalent graphs. This paper presents a novel technique for tensor graph superoptimization that employs equality saturation to apply all possible substitutions at once. We show that our approach can find optimized graphs with up to 16% speedup over state-of-the-art, while spending on average 48x less time optimizing.
This paper has not been read by Pith yet.
Forward citations
Cited by 7 Pith papers
-
Bonsai: Compiling Queries to Pruned Tree Traversals
Bonsai compiles queries to pruned tree traversals by deriving pruning conditions with extended symbolic interval analysis and fusing compound queries into single traversals.
-
Tensor Algebraic Property Skeletons: Amplifying Property-Based Testing for AI Compilers
Propilot instantiates 20 tensor-algebra property skeletons into 4,579 executable PBTs for TVM, cutting redundancy 49% and surfacing semantic and numerical errors.
-
Rewrite System Showdown: Stochastic Search vs. EqSat
The paper benchmarks equality saturation against stochastic search on five program optimization tasks to evaluate the effectiveness of e-graphs.
-
Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels
Nautilus auto-compiles math-like tensor descriptions into optimized GPU kernels, delivering up to 42% higher throughput than prior compilers on transformer models across NVIDIA GPUs.
-
Optimism in Equality Saturation
A new abstract interpretation algorithm enables sound optimistic analysis of e-graphs during equality saturation, unifying it with non-destructive rewriting and improving precision on cyclic SSA programs.
-
Rewrite System Showdown: Stochastic Search vs. EqSat
Empirical comparison of equality saturation versus stochastic search on five benchmarks to evaluate if e-graphs are superior for rewrite-based optimization.
-
LLM-Guided Strategy Synthesis for Scalable Equality Saturation
EggMind automates EqSat strategy synthesis via LLMs and EqSatL, cutting final cost 45.1% and peak RAM 69.1% versus full equality saturation on vectorization benchmarks while transferring to tensor compilers.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.