PaperlayerView exact revision on arXiv

Source-linked AI summary

Equality Saturation for Tensor Graph Superoptimization

Yichen Yang, Phitchaya Mangpo Phothilimtha, Yisu Remy Wang, Max Willsey, Sudip Roy, Jacques Pienaar

arXiv:2101.01332v2cs.AIcs.DC

TL;DR

Sequential graph rewriting is sensitive to rule order and may explore only a small part of the space of equivalent tensor graphs. TENSAT uses equality saturation, with extensions for tensor rewriting and extraction, and reports up to 16% speedup over state-of-the-art while spending an average of 48x less time optimizing.

  • Problem

    Existing tensor-graph rewriting and superoptimization apply substitutions sequentially, making search sensitive to rewrite order and often limiting exploration to a small part of the exponential space of equivalent graphs.

    Source: §1 · p. 1

  • Method

    TENSAT applies rewrites using equality saturation and extends exploration for complex, non-local rules while filtering invalid subgraphs for ILP-based extraction.

    Source: §1 · pp. 1–2

  • Results

    Up to 16% speedup over state-of-the-art and an average of 48x less optimization time are reported for TENSAT.

    Source: §8 · p. 10

  • Takeaways & Limitations

    The authors report globally optimal graphs that are fast enough for integration into the normal compilation flow for inference graphs.

    Source: §8 · p. 10

    For larger, more complex networks, exploration is limited to a certain number of multi-pattern rewrite iterations because the e-graph becomes too large for extraction.

    Source: §6.4 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Sequential graph rewriting can miss better tensor graphs because optimization depends on rewrite order. TENSAT uses equality saturation to explore equivalent graphs together, reporting faster optimized graphs and shorter optimization times than prior approaches.

1 INTRODUCTION

  • Up to 16% faster runtime than state-of-the-art, up to 68.9% faster than unoptimized graphs, and up to 300x less optimization time are reported for TENSAT.The comparison covers TASO and TENSAT on optimization time and optimized-graph runtime speedup over the original graphs.Source: Table 1 · p. 1, §1 · p. 2
  • Sequential heuristic and search-based methods are sensitive to rewrite order because rewrites can depend on or enable one another, creating a phase-ordering problem.The industry-standard approach also relies on manually curated rules and heuristic application order, which can yield sub-optimal results.Source: §1 · p. 1
  • TENSAT applies rewrites through equality saturation, using e-graphs to represent many equivalent graphs before extraction selects a low-cost graph.The paper extends exploration for complex, non-local rules and filters invalid subgraphs to help its ILP-based extraction find an optimal solution.Source: §1 · pp. 1–2

2 EQUALITY SATURATION BACKGROUND

Equality saturation addresses rewrite-choice problems by retaining equivalent terms in e-graphs rather than destructively replacing one term with another. It separates exploration from extraction, which chooses a low-cost equivalent term.

2.1 Term Rewriting

  • Term rewriting applies semantic-preserving rules to an expression by matching a pattern and substituting its corresponding right-hand side.Because rewriting is destructive, a poorly timed rewrite can hide a later optimization, as in the example where strength reduction prevents reducing (a × 2)/2 to a.Source: §2.1 · p. 2

2.2 E-graphs

  • An e-graph compactly encodes equivalent terms as e-classes containing e-nodes, whose children refer to other e-classes.A representative e-node can be chosen from each e-class to recover a represented term.Source: §2.2 · p. 2

2.3 Equality Saturation

  • Equality saturation grows an e-graph by applying rewrites until saturation or timeout, then extracts a lowest-cost represented term.The extracted term remains equivalent to the input when the rewrites are sound, and the exploration–extraction split avoids choosing rewrite order in advance.Source: §2.3 · p. 3
  • E-graphs can retain rewrites that may not help alone but enable other rules, and can prove some equalities directed rewriting cannot reach.The paper illustrates this additional proving power with rewrites over f(a, b), a, b, and c.Source: §2.3 · p. 3

3 TENSAT’S REPRESENTATIONS

TENSAT adapts a TASO-based representation of tensor computation graphs for equality saturation. It represents rewrite rules as symbolic patterns, including rules with multiple matched outputs.

3.1 Representing Tensor Computation Graphs

  • TENSAT represents tensor computation graphs as DAGs in a modified TASO-based representation, with operator nodes corresponding to output tensors and their inputs as children.Noop nodes combine final outputs to make the graph single-rooted without changing during exploration.Source: §3.1 · p. 3

3.2 Representing Rewrite Rules

  • TENSAT represents rewrite-rule source and target patterns as S-expressions with variables, and pairs corresponding outputs to state their equivalence.Single-output patterns are single-pattern rules; multi-output patterns are represented as lists of S-expressions and called multi-pattern rules.Source: §3.2 · pp. 3–4
  • Supported operators use tensor, string, integer, and tensor-tuple node types; integers encode parameters such as stride and axis, while strings encode variable-length parameters.The listed variable-length parameters include shapes and axis permutations.Source: Table 2 · p. 4
  • Figure 2 illustrates a multi-pattern rule that merges two matmuls sharing an input using concat and split, represented as S-expressions.Its identifiers prefixed with ? denote variables, and axis 1 is used for split and concat.Source: Figure 2 · p. 4, §3.2 · p. 4

4 TENSAT’S EXPLORATION PHASE

TENSAT’s exploration phase grows an e-graph by applying shape-compatible rewrites and adds a dedicated procedure for multi-pattern rules.

4 TENSAT’S EXPLORATION PHASE

  • Exploration starts from the input tensor graph and repeatedly adds rewrite targets and equivalences after checking target tensor shapes.It stops at saturation or a user-specified time, e-graph-size, or iteration limit.Source: §4 · p. 4
  • TENSAT searches canonicalized source patterns with the single-pattern routine, combines their matches, and applies only matches compatible at shared variables.Variable-renamed S-expressions share a canonical form for searching.Source: §4 · pp. 4–5
  • With N shared-input matmul nodes, a multi-pattern rule can create O(N^2) new nodes in one iteration and O(N^4) in the next.The passage describes this as double-exponential growth that can quickly explode the e-graph.Source: §4 · p. 5
  • Multi-pattern rules receive a separate iteration limit; afterward, TENSAT applies single-pattern rules until saturation or another user-specified limit.This bounds the period of multi-pattern application, which can rapidly grow the e-graph.Source: §4 · p. 5

5 TENSAT’S EXTRACTION PHASE

TENSAT extracts a minimum-cost acyclic graph from the e-graph using an ILP formulation and cycle filtering. Its filtering method reduces cycle-checking work by combining cached-descendant prechecks with post-processing.

5 TENSAT’S EXTRACTION PHASE

  • The cost model sums independently measured operator runtimes, a fit for GPUs that typically execute one operator at a time.Fused operators can count as a single operator.Source: §5 · p. 5

5.1 Extraction Algorithms

  • Greedy extraction chooses the lowest-cost subtree per e-class but can miss the minimum-cost graph when children share a subgraph.Ignoring shared subgraphs can make greedy extraction overestimate cost.Source: §5.1 · p. 5
  • ILP extraction guarantees a valid, cycle-free graph with the lowest cost; topological-order variables enforce the acyclic constraint.The formulation selects e-nodes subject to root and child-class constraints.Source: §5.1 · p. 6

5.2 Cycle Filtering

  • Efficient cycle filtering prechecks matches with cached descendants, then uses DFS post-processing to exclude the latest-added node in each missed cycle from extraction.Cycle constraints can slow the ILP solver, motivating filtering cycles during exploration.Source: §5.2 · pp. 6–7
  • Each pre-filter check takes constant time, and post-processing has worst-case complexity O(ncN); cycle counts are typically much smaller than rewrite-match counts, making filtering faster than the vanilla method.The post-processing bound is conservative because each DFS pass can find many cycles.Source: §5.2 · p. 7

6 EVALUATION

The evaluation compares TENSAT with TASO across seven inference models and examines optimization time, rewrite depth, and extraction design choices.

6.1 Experimental Setup

  • The evaluation uses seven inference models, TASO’s rewrite rules, and one NVIDIA Tesla T4 GPU; TENSAT’s main runs use kmulti = 1.TENSAT exploration is capped at 50,000 e-nodes or 15 iterations, with a one-hour ILP-solver timeout.Source: §6.1 · p. 7

6.3 Optimization Time

  • 9.5x to 379x faster than TASO’s total search time, TENSAT is also 1.8x to 260x faster than TASO’s best-found time.TASO’s best-found time is difficult, if not impossible, to attain in practice because the search cannot know when it has found its best graph.Source: Figure 5 · p. 8, §6.3 · p. 8

6.2 Program Speedup

  • TENSAT discovers better-optimized graphs than TASO in most benchmarks, with its largest speedup improvements on NasRNN and SqueezeNet.On Inception-v3, increasing kmulti from 1 to 2 yields better speedup than TASO while remaining 13.4× faster in search.Source: Figure 4 · p. 8, §6.2 · p. 8, Figure 6 · p. 8

6.4 Varying Iterations of Multi-Pattern Rewrites

  • Increasing kmulti improves optimized graphs for five benchmarks but expands the e-graph rapidly; for SqueezeNet, higher values can reduce speedup because the cost model diverges from runtime.Even at kmulti = 3, TENSAT on SqueezeNet still achieves better speedup than TASO.Source: §6.4 · pp. 8–9, Figure 7 · p. 9

6.5 Ablation Study

  • Greedy extraction fails to optimize BERT and NasNet-A because it chooses nodes independently and misses cases where multiple outputs must share a subgraph.The comparison covers BERT, NasRNN, and NasNet-A with exploration at kmulti = 1.Source: §6.5 · p. 9, Table 4 · p. 9
  • Removing cycle constraints makes ILP solving approximately 10x–1000x faster on larger e-graphs; with the constraints, solving times out at kmulti = 2 without finding a feasible solution.The authors identify satisfying the cycle constraints as the main difficulty for the ILP solver.Source: §6.5 · p. 9, Table 5 · p. 10
  • Efficient cycle filtering achieves up to 2000x speedup over vanilla filtering, enabling exploration of larger e-graphs.The comparison measures exploration-phase time as kmulti varies.Source: §6.5 · p. 9, Table 6 · p. 10

7 RELATED WORK

Prior work uses sequential substitutions for graph optimization and applies equality saturation in other domains. TENSAT extends tensor graph superoptimization to computation graphs with multi-pattern rewrites and cycle filtering for efficient ILP extraction.

7 RELATED WORK

  • Sampling-based search prunes redundant substitutions faster than TASO, but does not discover more optimized programs, unlike TENSAT.TASO uses backtracking with a hard threshold for substitutions that increase runtime.Source: §7 · p. 10
  • NeuRewriter uses reinforcement learning to iteratively choose rewrite rules and graph regions across domains including algebraic simplification, job scheduling, and vehicle routing.The authors describe it as complementary and suggest machine learning could select promising multi-pattern rules when saturation is not possible.Source: §7 · p. 10
  • Prior superoptimizers mostly target short low-level instruction sequences; Denali uses e-graphs and a constraint solver but sacrifices some optimality guarantee for scalability.The cited prior work does not focus on tensor graph superoptimization.Source: §7 · p. 10
  • Equality saturation has been applied to traditional compiler optimization and to CAD simplification, numerical accuracy, and code search.TENSAT applies it to tensor graph superoptimization.Source: §7 · p. 10
  • TENSAT optimizes tensor computation graphs, unlike Wang et al. (2020), who target linear algebra kernels with few simple operations.Its contributions include multi-pattern rewrites and cycle filtering for efficient ILP extraction.Source: §7 · p. 10

8 CONCLUSIONS

The conclusion presents equality saturation with multi-pattern rewrites and cycle-filtered extraction as a tensor graph optimization approach. It reports up to 16% speedup over state-of-the-art and average 48x less optimization time, and says the method is fast enough for inference compilation.

8 CONCLUSIONS

  • Up to 16% speedup over state-of-the-art and average 48x less optimization time are reported for TENSAT.The paper says its approach finds globally optimal graphs and is fast enough to integrate into the normal compilation flow for inference graphs.Source: §8 · p. 10
  • TENSAT extends equality saturation for tensor graph optimization with multi-pattern rewrite rules and an extraction algorithm using efficient cycle filtering.The extensions address making the approach work for the problem domain and extraction scalability.Source: §8 · p. 10
  • The authors report that TENSAT finds globally optimal graphs and is fast enough to integrate into the normal compilation flow for inference graphs.The conclusion frames this as a practical scope for the approach.Source: §8 · p. 10

A EXAMPLE PATTERNS OF REWRITE USED BY TENSAT

The examples show TENSAT rewrite patterns for BERT, NasNet-A, Inception-v3, and NasRNN. They include shared-input matrix multiplications or convolutions, and a NasNet-A pattern that combines four convolutions into two.

A EXAMPLE PATTERNS OF REWRITE USED BY TENSAT

  • BERT uses a pattern that generalizes to more than two matrix multiplications sharing a common input node.The figure identifies w1 and w2 as weight nodes.Source: Figure 8 · p. 13
  • NasNet-A and Inception-v3 use a pattern that generalizes to more than two convolutions sharing a common input node.The figure depicts an example of the pattern.Source: Figure 9 · p. 13
  • For NasNet-A, concatenating convolution weight kernels along output- and input-channel axes lets inference precompute the concatenations, effectively converting four convolutions into two.Kernel dimensions are ordered as output channels, input channels, height, width.Source: Figure 10 · p. 14
  • NasRNN uses a pattern that also generalizes to more than two matrix multiplications sharing a common input node.The figure gives an example pattern used in NasRNN.Source: Figure 11 · p. 14
Loading 2101.01332v2…