Equality Saturation for Tensor Graph Superoptimization

TL;DR

This paper introduces an equality saturation technique for tensor graph superoptimization, enabling simultaneous application of all substitutions, leading to more optimal graphs with significant speedups and reduced optimization time.

Contribution

The paper presents a novel equality saturation method for tensor graph superoptimization, overcoming limitations of sequential substitution approaches.

Findings

Achieves up to 16% speedup over state-of-the-art methods.

Reduces optimization time by an average of 48 times.

Effectively explores a larger space of equivalent graphs.

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.