A Learned Performance Model With Transfer Learning Across GPUs on Tensorized Instructions

Abstract

The training and inference efficiency of ever-larger deep neural networks highly rely on the performance of tensor operators on specific hardware accelerators. Therefore, a performance tuning framework with tensorized instruction compilation for automatic tensor generation is necessary for efficient deployment. These novel tensorized instruction, along with the emerging machine learning models, bring tremendous engineering challenges in compilation-based methods. They suffer from a large design space exploration with rough measurement accuracy and poor transferability among specialized instructions with certain hardware constraints. This paper presents a novel performance model for automatic code optimization with tensorized instruction. Central to the performance model is the assignment feature that not only clearly specifies the behaviour of instruction with computation and data movement abstraction, but also formally defines the matching problem from algorithm to tensorized instructions. Meanwhile, a simple yet efficient design with attention-inspired modules to accurately predict the performance of optimized tensor program by capturing global and long-range dependencies within a complete scheduling space. Compared with state-of-the-arts, our performance model can predict the optimal implementation of code configurations with tensorized instruction to reduce inference latency and search time by up to 1.21× and 3.41× on modern DNN benchmarks. Furthermore, with pre-trained parameters, our performance can quickly adapt to different workloads and platforms on tensorized instruction via transfer learning.