Automatically Generating ML Compiler Backends from Tensor Accelerator ISA Descriptions
This program is tentative and subject to change.
Machine learning (ML) compilers play a key role in enabling high-performance implementations of ML workloads. These compilers use existing CPU and GPU backends to generate device-specific code. In recent years, many tensor accelerators (or AI accelerators) have been designed to further accelerate these workloads, with commercial products like AWS Trainium publicly available. However, compared to commodity hardware, a majority of tensor accelerators do not have mature ML compiler backends with robust code generation support. Moreover, tensor accelerator designs are subject to fast iteration cycles, making it difficult to manually develop and maintain ML compiler backends. Therefore, to enable faster integration of novel tensor accelerator designs in ML infrastructure, we need to make the compiler backend construction process more agile.
In this paper, we introduce ACT, a compiler backend generator that automatically generates compiler backends for tensor accelerators, given just the instruction set architecture (ISA) descriptions. These backends are integrated with XLA, a production ML compiler. ACT uses a novel ISA-parameterized compilation algorithm to generate a compiler backend with an equality-saturation-based instruction selection phase and a constraint-programming-based memory allocation phase. We generated compiler backends for 6 accelerator platforms from industry (e.g., AWS Trainium, Intel AMX) and academia (e.g., Gemmini). We showed that these generated backends match or outperform commercial compiler backends and expert-written kernel libraries, while maintaining low compilation overheads. Notably, ACT-generated backend for AWS NKI ISA improved the code generation coverage for AWS Trainium by 2.3x compared with AWS’s production compiler, neuronx-cc.
ACT is part of a larger open-source ecosystem, built around our ISA description language TAIDL, that automatically generates essential software tools, such as test oracles and compiler backends, from ISA descriptions of tensor accelerators. Our tooling has been adopted by multiple academic and industry teams designing novel tensor accelerators. The ecosystem is available at https://github.com/act-compiler/act.
This program is tentative and subject to change.
Tue 6 OctDisplayed time zone: Pacific Time (US & Canada) change
15:30 - 17:00 | Compiler Optimization and Code GenerationOOPSLA at East Hall 1 Chair(s): Kirshanthan Sundararajah Virginia Tech | ||
15:30 18mTalk | Class-Dictionary Specialization with Rank-2 Polymorphic Functions OOPSLA Link to publication DOI Pre-print | ||
15:48 18mTalk | Automatic Propagation of Profile Information through the Optimization Pipeline OOPSLA Elisa Frohlich Federal University of Minas Gerais, Angelica Moreira Microsoft Research, Fernando Magno Quintão Pereira Federal University of Minas Gerais DOI | ||
16:06 18mTalk | Automatically Generating ML Compiler Backends from Tensor Accelerator ISA Descriptions OOPSLA Devansh Jain University of Illinois at Urbana-Champaign, Akash Pardeshi University of Illinois at Urbana-Champaign, Marco Frigo University of Illinois at Urbana-Champaign, Kaustubh Khulbe University of Illinois at Urbana-Champaign, Krut Patel NVIDIA, Saatvik Lochan University of Illinois at Urbana-Champaign, Jai Arora University of Illinois at Urbana-Champaign, Charith Mendis University of Illinois at Urbana-Champaign DOI Pre-print | ||
16:24 18mTalk | Symbolic Basic Block Profiling for Machine Learning Kernels OOPSLA Jingyu Qiu University of Rochester, Rongcui Dong University of Rochester, Sreepathi Pai University of Rochester DOI | ||
16:42 18mTalk | Filtr: Compiling Bioinformatics Recurrences OOPSLA Bala Vinaithirthan Stanford University, Shiv Sundram Stanford University, Sneha Goenka Princeton University, Fredrik Kjolstad Stanford University DOI | ||