Publications

2026

1. OOPSLA

   Automatically Generating ML Compiler Backends from Tensor Accelerator ISA Descriptions

   Devansh Jain, Akash Pardeshi, Marco Frigo, Kaustubh Khulbe, Krut Patel, Saatvik Lochan, Jai Arora, and Charith Mendis

   Proc. ACM Program. Lang., Oct 2026

   

   Abs DOI arXiv Bib PDF Code

   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{10.1145/3839457,
     author = {Jain, Devansh and Pardeshi, Akash and Frigo, Marco and Khulbe, Kaustubh and Patel, Krut and Lochan, Saatvik and Arora, Jai and Mendis, Charith},
     title = {Automatically Generating ML Compiler Backends from Tensor Accelerator ISA Descriptions},
     year = {2026},
     issue_date = {October 2026},
     publisher = {Association for Computing Machinery},
     address = {New York, NY, USA},
     volume = {10},
     number = {OOPSLA2},
     url = {https://doi.org/10.1145/3839457},
     doi = {10.1145/3839457},
     journal = {Proc. ACM Program. Lang.},
     month = oct,
     articleno = {325},
     numpages = {32},
     keywords = {AI Accelerators, Automated Compiler Construction, ML Compilers},
   }

2025

1. MICRO

   TAIDL: Tensor Accelerator ISA Definition Language with Auto-generation of Scalable Test Oracles

   Devansh Jain, Marco Frigo, Jai Arora, Akash Pardeshi, Zhihao Wang, Krut Patel, and Charith Mendis

   In Proceedings of the 58th IEEE/ACM International Symposium on Microarchitecture (MICRO), Oct 2025

   

   Abs DOI Bib PDF

   With the increasing importance of deep learning workloads, many hardware accelerators have been proposed in both academia and industry. However, software tooling for the vast majority of them does not exist compared to the software ecosystem and innovations proposed for established platforms such as CPUs and GPUs. We observed that the lack of well-defined hardware-software interfaces and correctness testing tools like fast and scalable test oracles (also known as functional simulators) act as significant barriers to adopting these emerging accelerators in the software community. These interfaces and tools are essential in building software such as retargetable compilers and optimized kernels. To bridge these gaps, we first present TAIDL, an instruction specification language that provides novel constructs to describe the instruction set architectures (ISAs) of tensor accelerators. Next, given ISA definitions in TAIDL, we introduce techniques to automatically generate fast and scalable test oracles for diverse sets of accelerators, which are needed for testing software correctness of code that targets pre-silicon hardware designs. Automated generation of such tools reduces the burden on hardware architects and the repeated development efforts required across different accelerator platforms. Further, our techniques allow us to execute these simulators on GPUs, leading to highly scalable simulations. To demonstrate the expressivity of TAIDL, we instantiated several tensor accelerator ISAs with different compute capabilities and memory hierarchies. Further, we show that test oracles generated using TAIDL definitions are orders of magnitude faster and more scalable than existing instruction-level functional simulators, making them suitable for integration into software development cycles. TAIDL is available at https://github.com/act-compiler/taidl.

   @inproceedings{10.1145/3725843.3756075,
     author = {Jain, Devansh and Frigo, Marco and Arora, Jai and Pardeshi, Akash and Wang, Zhihao and Patel, Krut and Mendis, Charith},
     title = {TAIDL: Tensor Accelerator ISA Definition Language with Auto-generation of Scalable Test Oracles},
     year = {2025},
     isbn = {9798400715730},
     publisher = {Association for Computing Machinery},
     address = {New York, NY, USA},
     url = {https://doi.org/10.1145/3725843.3756075},
     doi = {10.1145/3725843.3756075},
     booktitle = {Proceedings of the 58th IEEE/ACM International Symposium on Microarchitecture (MICRO)},
     month = oct,
     pages = {1316-1333},
     numpages = {18},
     series = {MICRO '25},
   }

2. POPL

   TensorRight: Automated Verification of Tensor Graph Rewrites

   Jai Arora, Sirui Lu, Devansh Jain, Tianfan Xu, Farzin Houshmand, Phitchaya Mangpo Phothilimthana, Mohsen Lesani, Praveen Narayanan, Karthik Srinivasa Murthy, Rastislav Bodik, Amit Sabne, and Charith Mendis

   Proc. ACM Program. Lang. (POPL), Jan 2025

   🏆 Distinguished Paper Award
   

   Abs DOI arXiv Bib PDF Supp Code

   Tensor compilers, essential for generating efficient code for deep learning models across various applications, employ tensor graph rewrites as one of the key optimizations. These rewrites optimize tensor computational graphs with the expectation of preserving semantics for tensors of arbitrary rank and size. Despite this expectation, to the best of our knowledge, there does not exist a fully automated verification system to prove the soundness of these rewrites for tensors of arbitrary rank and size. Previous works, while successful in verifying rewrites with tensors of concrete rank, do not provide guarantees in the unbounded setting. To fill this gap, we introduce TensorRight, the first automatic verification system that can verify tensor graph rewrites for input tensors of arbitrary rank and size. We introduce a core language, TensorRight DSL, to represent rewrite rules using a novel axis definition, called aggregated-axis, which allows us to reason about an unbounded number of axes. We achieve unbounded verification by proving that there exists a bound on tensor ranks, under which bounded verification of all instances implies the correctness of the rewrite rule in the unbounded setting. We derive an algorithm to compute this rank using the denotational semantics of TensorRight DSL. TensorRight employs this algorithm to generate a finite number of bounded-verification proof obligations, which are then dispatched to an SMT solver using symbolic execution to automatically verify the correctness of the rewrite rules. We evaluate TensorRight’s verification capabilities by implementing rewrite rules present in XLA’s algebraic simplifier. The results demonstrate that TensorRight can prove the correctness of 115 out of 175 rules in their full generality, while the closest automatic, bounded-verification system can express only 18 of these rules.

   @article{10.1145/3704865,
     author = {Arora, Jai and Lu, Sirui and Jain, Devansh and Xu, Tianfan and Houshmand, Farzin and Phothilimthana, Phitchaya Mangpo and Lesani, Mohsen and Narayanan, Praveen and Murthy, Karthik Srinivasa and Bodik, Rastislav and Sabne, Amit and Mendis, Charith},
     title = {TensorRight: Automated Verification of Tensor Graph Rewrites},
     year = {2025},
     issue_date = {January 2025},
     publisher = {Association for Computing Machinery},
     address = {New York, NY, USA},
     volume = {9},
     number = {POPL},
     url = {https://doi.org/10.1145/3704865},
     doi = {10.1145/3704865},
     journal = {Proc. ACM Program. Lang. (POPL)},
     month = jan,
     articleno = {29},
     numpages = {32},
     keywords = {Denotational Semantics, Tensor Compilers, Unbounded Verification},
   }