一句话：提出 TensorIR，将张量计算作为一等公民的编译器抽象，支持自动张量化优化，性能媲美手工优化系统。
在多种设备上部署深度学习模型面临硬件专用加速原语多样、工程复杂度高等问题。本文提出 TensorIR，一种用于优化带张量计算原语程序的编译器抽象。它泛化现有机器学习编译器的循环嵌套表示，将张量计算提升为一等公民，从而统一描述计算与硬件加速原语。作者在 TensorIR 之上构建端到端框架，针对给定张量计算原语自动优化深度学习模型。实验表明，TensorIR 编译能自动调用目标硬件后端的张量计算原语，并在多个平台上达到与当前最先进手工优化系统相竞争的性能。
Deploying deep learning models on various devices has become an important topic. The wave of hardware specialization brings a diverse set of acceleration primitives for multi-dimensional ten- sor computations. These new acceleration primitives, along with the emerging machine learning models, bring tremendous engineering challenges. In this paper, we present TensorIR, a compiler abstraction for optimizing programs with these tensor computation primitives. TensorIR generalizes the loop nest representation used in existing machine learning compilers to bring tensor computation as the first-class citizen. Finally, we build an end-to-end framework on top of our abstraction to automatically optimize deep learning models for given tensor computation primitives. Experimental results show that TensorIR compilation automatically uses the tensor computation primitives for given hardware backends and delivers performance that is competitive to state-of-art hand-optimized systems across platforms.
@inproceedings{paperbot9992,
title = {TensorIR: An Abstraction for Automatic Tensorized Program Optimization},
author = {Siyuan Feng and Bohan Hou and Hongyi Jin and Wuwei Lin and Junru Shao and Ruihang Lai and Zihao Ye and Lianmin Zheng and Cody Hao Yu and Yong Yu and Tianqi Chen},
booktitle = {ASPLOS},
year = {2023},
doi = {10.1145/3575693.3576933}
}