ATF ormer: A Learned Performance Model with Transfer Learning Across Devices for Deep Learning Tensor Programs (2023)

Abstract: The training and inference efficiency of ever-larger deep neural networks highly rely on the performance of tensor operators on specific hardware platforms. Therefore, a compilation-based optimization flow with automatic tensor generation and parameter tuning is necessary for efficient model deployment. While compilation-based methods with performance models can provide dynamic and suitable code optimization, they suffer from a large design space exploration with rough measurement accuracy and poor transferability among different hardware platforms. This paper presents ATFormer, a simple yet efficient design with attention-inspired modules to accurately predict the performance of optimized operators by capturing global and long-range dependencies within a complete scheduling space. Compared with state-of-the-arts, ATFormer can predict the optimal implementation of tensor operators to reduce inference time with minimal effort on modern DNN benchmarks. Furthermore, ATFormer with pre-trained parameters can quickly adapt to different workloads and hardware via transfer learning.

## introduction recently, there has been a significant improvement in model performance for deep neural networks (dnns) @xcite @xcite @xcite @xcite . however, this progress has been accompanied by a significant increase in the number of operators and, consequently, the computational complexity of dnns. as a result, it has become increasingly challenging to efficiently deploy dnns with optimized tensor programs on certain hardware accelerators like cpus, gpus and tpus @xcite . to overcome the limitations, mainstream searchbased tensor compilers @xcite @xcite @xcite are developed. these compilers automatically search for the optimal deployment configuration of each operator on increasingly heterogeneous platforms. conducting on-device measurements is extremely time-consuming, making it impossible to place all the generated tensor programs on the target platform for measurement during the compilation process. therefore, the prediction via an optimal cost model is crucial in reducing the time-consuming measurements during the compilation which can significantly improve search efficiency and quality. nevertheless, the existing cost models are capable of selecting nearly optimal configurations but suffer from excessively long optimization time. these long optimization times not only impede the deployment period but also raise concerns about the practicality of search-based compilers. furthermore, statistic cost models trained on one hardware platform exhibit significant performance degradation on different hardware, making them unusable across different platforms. it is noteworthy that the execution times of tensor programs can vary signif-icantly on different platforms due to domain gaps, making it challenging to deploy optimized models on multiple platforms. this is further compounded by the significant differences in the features extracted from various platforms. even when extracted on gpus, the feature's stability and performance cannot be guaranteed across different gpu architectures such as volta, turing, and ampere. therefore, additional engineering efforts are necessary to account for the differences in hardware architectures, resulting in a laborious and cumbersome feature extraction process. to address these challenges, we propose a powerful yet simple approach that uses attention-inspired blocks to enhance the performance of cost models. these blocks can capture global and long-range dependencies among tensor program statements. additionally, transferable features with pre-trained parameters are used to expedite search convergence across different hardware platforms. these techniques can be easily incorporated into existing search algorithms and improve efficiency in an endto-end fashion. our design, atformer, consistently outperforms popular dnn benchmarks, including small and large-scale models. furthermore, our techniques enable cross-platform transfer learning, resulting in more efficient deployment. the main contributions of this paper are the following: (i) we highlight the limitations of current auto-tuning frameworks. existing tree-based performance models are insufficient for evaluating inference in a large search space and transferable knowledge is difficult to acquire across different platforms. (ii) a simple yet efficient design that utilizes attention-based blocks to explore the correlation between all innermost non-loop statements in a full tensor program, resulting in accurate prediction. (iii) our approach enables rapid adaptation of performance tuning across various gpu platforms using pre-trained parameters on static datasets, not only in cross-operator but also crossplatform scenarios. comprehensive experiments on modern dnn benchmarks and the large-scale tenset @xcite demonstrate the consistent and superior performance of our method. tree-based performance model. decision trees are frequently used in classification and regression problems. to enhance their performance, an ensemble learning approach is typically employed to reduce variance. xgboost @xcite and lightgbm are powerful feature models in sequence modeling tasks. to achieve accurate prediction, a number of works, including @xcite @xcite @xcite @xcite , use xgboost as the performance model during the tuning. autotvm extracts domain-specific features from a provided low-level abstract syntax tree (ast). during optimization, these features, which include loop structure information and generic annotations, are explored. moreover, treegru @xcite recursively encodes a low-level ast into an embedding vector, which is mapped to a final predicted score within a fully-connected layer to enhance performance. halide @xcite builds regression models with hardware-specific features for auto-scheduling. tabnet (arık and pfister, 2020) uses sequential attention to select the most salient features to reason at each decision via a deep tabular architecture. ## conclusion this paper introduces atformer, a novel and effective design for optimizing tensor programs. at-former employs hierarchical features with varying levels of granularity to model the end-to-end compilation. moreover, self-attention blocks are utilized to explore global dependencies of a complete tensor program for high-quality evaluation. through transfer learning, atformer achieves faster-converged latency and superior transferability across different hardware platforms, outperforming previous state-of-the-art benchmarks. limitations. we plan to do the transfer learning from gpus to cpus and explore the potential of combining with post-training quantization or pruning to efficiently deploy models. additionally, we will explore more universal and efficient methods for optimizing tensor programs with atformer. this includes leveraging hardware features to optimize performance on domain-specific accelerators, such as nvidia's tensor cores.