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...

Ausführliche Beschreibung

Gespeichert in:
Bibliographische Detailangaben
Veröffentlicht in:2021 30th International Conference on Parallel Architectures and Compilation Techniques (PACT) S. 1 - 16
Hauptverfasser: 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: Tagungsbericht
Sprache:Englisch
Veröffentlicht: IEEE 01.09.2021
Schlagworte:
autotuning
compiler
Internet
Layout
Machine learning
Production
Search problems
Tensors
Training
Online-Zugang:Volltext
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
Beschreibung
Zusammenfassung: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