TASO: optimizing deep learning computation with automatic generation of graph substitutions

By Oded Padon, James J. Thomas, Todd Warszawski, Matei A. Zaharia, A. Aiken. Symposium on Operating Systems Principles, 2019

371 citations recorded by the source corpus.

Existing deep neural network (DNN) frameworks optimize the computation graph of a DNN by applying graph transformations manually designed by human experts. This approach misses possible graph optimizations and is difficult to scale, as new DNN operators are introduced on a regular basis. We propose TASO, the first DNN computation graph optimizer that automatically generates graph substitutions. TASO takes as input a list of operator specifications and generates candidate substitutions using the given operators as basic building blocks. All generated substitutions are formally verified against the operator specifications using an automated theorem prover. To optimize a given DNN computation graph, TASO performs a cost-based backtracking search, applying the substitutions to find an optimized graph, which can be directly used by existing DNN frameworks. Our evaluation on five real-world DNN architectures shows that TASO outperforms existing DNN frameworks by up to 2.8X, while requiring significantly less human effort. For example, TensorFlow currently contains approximately 53,000 lines of manual optimization rules, while the operator specifications needed by TASO are only 1,400 lines of code.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/80b362efee95c1759c6dab9219eb77ca3ee44475

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

Citations in the corpus

This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (positive): “Compiler techniques have been introduced to optimize the execution of DL models [13,24,34,51,52,54,62].”
A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters (2020, background): “, Tensor Fusion [14] and graph substitution [37].”
Ansor : Generating High-Performance Tensor Programs for Deep Learning (2020, background): “The common optimizations at graph level include layout optimizations [32], operator fusion [11, 38, 60], constant folding [42], auto-batching [33], automatic generation of graph substitution [29] and so forth.”
egg: Fast and extensible equality saturation (2020, methodology): “To demonstrate this, we replicated a portion of the recent TASO paper [Jia et al. 2019], which optimizes deep learning models.”
Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, methodology): “Beyond these, TASO (Jia et al., 2019a) targets similar optimizations to ours by using graph substitutions, but considers only inference and does not exhaustively explore the search space.”
DistIR: An Intermediate Representation for Optimizing Distributed Neural Networks (2021, methodology): “Many recent works propose algorithms for automatic distribution, such as MCMC search [17], integer and dynamic programming [24, 33], reinforcement learning [23, 34], and custom algorithms [16, 25].”