TensorIR: An Abstraction for Automatic Tensorized Program Optimization
This addresses the engineering problem of efficiently deploying models on specialized hardware for developers and researchers, though it appears incremental by building on existing compiler representations.
The paper tackles the challenge of optimizing deep learning models for diverse hardware acceleration primitives by introducing TensorIR, a compiler abstraction that makes tensor computations first-class citizens, and demonstrates that it achieves performance competitive with state-of-the-art hand-optimized systems across platforms.
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 tensor 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.