A Flexible Approach to Autotuning Multi-Pass Machine Learning Compilers

Search-based techniques have been demonstrated effective in solving complex optimization problems that arise in domain-specific compilers for machine learning (ML). Unfortunately, deploying such techniques in production compilers is impeded by two limitations. First, prior works require factorizatio...

Full description

Saved in:
Bibliographic Details
Published in2021 30th International Conference on Parallel Architectures and Compilation Techniques (PACT) pp. 1 - 16
Main Authors Phothilimthana, Phitchaya Mangpo, Sabne, Amit, Sarda, Nikhil, Murthy, Karthik Srinivasa, Zhou, Yanqi, Angermueller, Christof, Burrows, Mike, Roy, Sudip, Mandke, Ketan, Farahani, Rezsa, Wang, Yu Emma, Ilbeyi, Berkin, Hechtman, Blake, Roune, Bjarke, Wang, Shen, Xu, Yuanzhong, Kaufman, Samuel J.
Format Conference Proceeding
LanguageEnglish
Published IEEE 01.09.2021
Subjects
autotuning
compiler
Internet
Layout
Machine learning
Production
Search problems
Tensors
Training
Online AccessGet full text

Cover

More Information
Summary:Search-based techniques have been demonstrated effective in solving complex optimization problems that arise in domain-specific compilers for machine learning (ML). Unfortunately, deploying such techniques in production compilers is impeded by two limitations. First, prior works require factorization of a computation graph into smaller subgraphs over which search is applied. This decomposition is not only non-trivial but also significantly limits the scope of optimization. Second, prior works require search to be applied in a single stage in the compilation flow, which does not fit with the multi-stage layered architecture of most production ML compilers. This paper presents Xtat, an autotuner for production ML compilers that can tune both graph-level and subgraph-level optimizations across multiple compilation stages. Xtat applies Xtat-M, a flexible search methodology that defines a search formulation for joint optimizations by accurately modeling the interactions between different compiler passes. Xtat tunes tensor layouts, operator fusion decisions, tile sizes, and code generation parameters in XLA, a production ML compiler, using various search strategies. In an evaluation across 150 ML training and inference models on Tensor Processing Units (TPUs) at Google, Xtat offers up to 2.4x and an average 5% execution time speedup over the heavily-optimized XLA compiler.
DOI:10.1109/PACT52795.2021.00008