一句话:提出 ACT,仅凭张量加速器 ISA 描述自动生成集成 XLA 的 ML 编译器后端,性能媲美或超越商用后端与手写库。
ML 编译器通常依赖成熟的 CPU/GPU 后端,而众多张量加速器缺乏稳健的编译器后端支持,且加速器设计迭代快,手工开发与维护后端成本高。本文提出 ACT,一种编译器后端生成器:输入仅需张量加速器的指令集架构(ISA)描述,即可自动生成与 XLA 集成的 ML 编译器后端。ACT 采用 ISA 参数化的编译算法,包含基于等式饱和(equality saturation)的指令选择阶段和基于约束编程的内存分配阶段。作者为 6 个来自工业界(如 AWS Trainium、Intel AMX)与学术界(如 Gemmini)的加速器平台生成后端,实验表明其性能匹配或优于商用编译器后端与专家手写 kernel 库,同时保持较低编译开销;面向 AWS NKI ISA 生成的后端将 Trainium 的代码生成覆盖率较 AWS 生产编译器 neuronx-cc 提升 2.3 倍。该工作基于开源 ISA 描述语言 TAIDL,构成自动生成测试预言与编译器后端等工具的开源生态,并已被多个学术与工业团队采用。
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.
@article{paperbot12168,
title = {Automatically Generating ML Compiler Backends from Tensor Accelerator ISA Descriptions},
author = {Devansh Jain and Akash Pardeshi and Marco Frigo and Kaustubh Khulbe and Krut Patel and Saatvik Lochan and Jai Arora and Charith Mendis},
journal = {Proceedings of the ACM on Programming Languages},
volume = {10},
number = {OOPSLA2},
year = {2026},
doi = {10.1145/3839457}
}