ATFormer: A Learned Performance Model with Transfer Learning Across Devices for Deep Learning Tensor Programs
Authors: Yang Bai
Wenqian Zhao
Shuo Yin
Zixiao Wang
Bei Yu
Venue: EMNLP 2023 Main (2023)
Keywords: Tensor Program; Performance Model; Efficient Transfer Learning; NLP Application; Model Deployment

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.

ATFormer: A Learned Performance Model with Transfer Learning Across Devices for Deep Learning Tensor Programs
Yang Bai, Wenqian Zhao, Shuo Yin, Zixiao Wang, Department of Computer Science and Engineering The Chinese University of Hong Kong, Hong Kong SAR [ybai, wqzhao, syin22, zxwang22, byueecse. cuhk. edu. hk
Bei Yu
Abstract
Imput
DL framework fromtends
High-level structure
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 ten- sor generation and parameter tuning is neces- sary for efficient model deployment. While compilation- based methods with performance models can provide dynamic and suitable code optimization, they suffer from large design space exploration with rough measurement ac- curacy and poor transferability among differ- ent hardware platforms. This paper presents ATFormer, simple yet efficient design with attention-inspan iintiiion-nsieeed modules to accurately pre- dict the performance of optimized operators by capturing global and long-range range dependen- cies within complete scheduling space. Com- pared with late-of-theearte ATFormer can predict the optimal implementation of tensor operators reduce inference time with mini- mal effort modern DNN benchmarks. Fur- thermore, ATFormer with pre-trained parame- can quickly adapt to different workloads and hardware via transfer learning.
Conv [a]
mxnet
[b]
Conv [b]
Low-level
detail
Conv [c
Matmal lc]
Computation
Gruph
Output
ATFommer
Cost Model
Search Algorithu
Cool geneatton buckends
Online DataSet
Offinne
Search Space
Best Candidate
DataSet
Hardware
Full Tensor Programs
Auto-Tuning
Figure 1: The overview of search-based framework with computation graph, cost model, and search space.
automatically search for the optimal deployment configuration of each operator increasingly het- erogeneous platforms. Conducting on-device mea- surements is extremely time-consuming, making impossible to place all the generated tensor pro- grams on the target platform for measurement dur- ing the compilation process. Therefore, the pre- diction via an optimal cost model is crucial re- ducing the time-consuming measurements during the compilation which can significantly improve search efficiency and quality.
1 Introduction
Recently, there has been significant improvement in model performance for deep neural networks (DNNs) (He et al., 2016; Sandler et al., 2018; Shan al., 2021: Devlin al., 2019; Wu al., 2019; Biten et al., 2019; Bello et al., 2019). However, this progress has been accompanied by signifi- cant increase the number of operators and, con- sequently, the computational complexity of DNNs. As result, it has become increasingly challeng- ing to efficiently deploy DNNs with optimized ten- sor programs certain hardware accelerators like CPUs, GPUs and TPUs (Jouppi al., 2017).
Nevertheless, the existing cost models are ca- pable 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 barch-based compilers. Further- more, statistic cost models trained on one hardware platform exhibit significant performance degrada- tion on different hardware, making them unusable across different platforms. It is noteworthy that the execution times of tensor programs can vary signif-
To overcome the limitations, mainstream search- based tensor compilers (Chen et al., 2018aa 8a; Zheng et al., 2020; Bai et al., 2021; Li et al., 2020; Fe- gade et al., 2021) are developed. These compilers

---

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 ex- tracted from various platforms. Even when ex- tracted on GPUs, the feature's 's stability and perfor- mance cannot be guaranteed across different GPU architectures such as Volta, Turing, and Ampere. Therefore, additional engineering efforts are nec- essary to account for the differences hardware architectures, resulting laborious and cumber- some feature exiraction process.
and TACO (Kjolstad al., 2017), has progressed rapidly. These optimization schemes typically con- sist of two parts: DL framework frontends and code generation backends, as illustrated in Figure 1. The frontend converts an input model into a high-level graph-based intermediate representation (IR) and applies anget-independent optimizations, such operator fusion and data layout transformation. In the backend, target-dependent optimization passes, along with hardware features, further optimize the final performance. TVM (Chen et al., 201 8a) is a slate-of-the-- art scarch-based tensor compiler that is widely used in academia and industry. Its auto- tuning aims to achieve performance comparable to hand-tailored libraries and has achieved promis- ing results. TVM has two versions of auto-tuning: AuM (Chen al., 201 2018c) and Ansor (Zheng et al., 2020). While AutoTVM is a semi-automated framework that requires pre-defined manual tem- plates, Ansor more advanced and fully auto- mated. However, both frameworks need to collect data on-the-fly during the search, resulting in an extremely long compilation time.
To address these challenges, we propose power- ful yet simple approach that uses inientiioninsspire 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 tech- niques can be easily incorporated into existing search algorithms and improve efficiency in an end- to-end fashion. Our design, ATFormer, consistently outperforms popular DNN benchmarks, including small and large-scale models. Furthermore, our techniques enable cross-platfone transfer learning, resulting in more efficient deployment.
Tree-based Performance Model. Decision trees are frequently used in classification and regression problems. To enhance their performance, an en- semble learning approach is typically employed to reduce variance. XGBoost (Chen and Guestrin, 2016a) and LightGBM are powerful feature models in sequence modeling tasks. To achieve accurate prediction, a number of works, including (Chen et al., 201 8c; Zheng et al., 2020; Ahn et al., 2020; Gao et al., 2021; Bai et al., 2021, 2023; Huang et al., 2023; Zhao et al., 2023), use XGBoost as the performance model during the tuning. AVOoTVM extracts domain-specift features from provided low-level abstract syntax tree (AST). During opti- mization, these features, which include loop struc- ture information and generic annotations, are ex- plored. Moreover, TreeGRU (Tai et al., 2015) recur- sively encodes a low-level AST into an embedding vector, which is mapped to a final predicted score within fully-connected layer to enhance perfor- mance. Halide (Adams et al., 2019) builds regres- sion models with hardware-specific features for auto-scheduling. TabNet (Ank and Pfister, 2020) uses sequential attention select the most salient features to reason at each decision via a deep tabu- lar architecture.
The main contributions of this paper are the fol- lowing: (i) We highlight the limitations of current auto-tuning frameworks. Existing tree-based per- formance models are insufficient for evaluating in- ference a large search space and transferable knowledge is difficult to acquire across different platforms. (ii) A simple yet efficient design that blocks to explore the corre- lation between all innermost loop statements in full tensor program, resulting in accurate pre- diction. (iii) Our approach enables rapid adapta- tion of performance tuning across various GPU platforms using pre-trained parameters on static datasets, not only in cross-operator but also cross- platform scenarios. Comprehensive experiments on modern DNN benchmarks and the large-scale TenSet (Zheng et al., 2021) demonstrate the consis- tent and superior performance of our method.
utilizes
2 Background and Related Work
Deep Learning Compiler. Recently, the devel- opment of compiler-based optimization frame- works, such as Halide (Adams et al., 2019), TVM (Chen et al., 2018b), XLA (Sabne, 2020),
DNN-based Performance Model. In contrast, some recent approaches aim to reduce the impact of

---

search algorithms on final performance by utilizing more robust and powerful cost models. (Kaufman et al., 2020) and (Sun et al., 2022) employ graph neural networks to predict the latency of DNNs on TPUs. (Steiner et al., 2021) formulates the tuning process as a deterministic Markov Decision Pro- cess (Xiang et al., 2015) and solves it by learning approximation of the value function. Tiramisu (Baghdadi al., 2019) manually extracts 2534 fea- tures from the structure of AST, and forwards the AST as a computation stream propagate features during the training. These models are trained effec- tively on a dataset with only a few thousand sched- ules using the hardware-dependent features crafted by heavy feature engineering techniques. However, complex feature engineering can become problem- atic in such cases. As hardware-specific features are difficult to transfer a new platform, learned performance model trained on one hardware plat- form typically performs poorly on another. This leads to an issue we call cross hardwwh unavail- ability. Additionally, this approach cannot keep pace with the rapid development of new hardware, which further exacerbates the problem.
Algorithm 1 Search-based Framework
Input: Search space Φ1, dy with operator or and setting k. Output: Tensor program p' with best configuration c².
1: while n'Trials < eachSubgraph Trials do
2: GS1 + SeneratetighhSetchhee ,0,
3:
CS2 <- Sampling (GS) , d2. o, o.k);
4:
P + EvolutionSearch(C
GS.):
5:
for P P do
6:
c e- f(0((1 O2 a. k]);
7:
end for
8:
nTrials <- nTrials + batchSize:
9:
end while
10: c" <- best lensor program configurations:
are dependent the hardware platform. Each ten- sor program can be considered as a candidate in the search space. We define the hierarchical search space $1,2, which decouples high-level structures Φ1 from low-level details Φ2, allowing for the ef- ficient exploration of potential tensor candidates during the tuming process.
Here, we can transform a tuning problem into an optimization problem that explores the potential tensor programs in a hierarchical search space.
Problem Given code generation function ö, high-level structure generation parameters é1, low- level detail sampling parameters 122, computation- intensive operator a and operator setting & (e.g.. kernel size), our goal is to use Φ1.2 to build a hier- archical search space and generate tensor program to achieve the optimal prediction score y² a specific hardware platform H.
3 Methodology
3.1 Problem Formulation
We describe DNN model as a computation graph and then define some important terminologies.
Definition 1 (Subgraph). Computation Graph G is partitioned into set of subgraphs S based on the graph-level optimizer (Roesch et al., 2018).
$ = arg max y, y fa(o o2)a,
Each search task is extracted from an indepen- dent subgraph S; on a specific hardware platform H. Thus, we define search task Q as follows:
(2)
The cost model f predicts score y of the tensor program p. The accuracy of the cost model f is crucial in finding ideal optimization configuration.
QU(SIG) { , Pi - -.. {s
(1)
where is the number of subgraphs in G. Note that each subgraph S, contains a computation-intensiv operator a and represent the i- th search task in G. Each sub- graph S; has own search space, which is deter- mined by the input and output shapes, data preci- sions, memory layout, and the hardware platform. The search space is usually large enough to cover almost all kinds of tensor candidates.
3.2 Performance Model
The process of optimization using our design is outlined in Algorithm 1. The input is set of to- be-optimized operators or subgraphs with different configurations. To implement our workflow, three functions are defined: SenneratetighSkkttee Sampling() (), and EvolutionSearch() as shown takes
S,. Therefore, use ((S.IC)
in Algorithm Φ1. o, and k input and returns the high-level generation sketch GS₁ as output. Sampling() takes GS}, d2, o, and h as input and returns the low-level annotation samples GS2 as output.
Definition (Hierarchical Search Space). tensor program, denoted by represents an implementa- tion of the subgraph using low-level primitives that

---

Hlgh-level Structure:
Compuiiiin Layer
Ergieioo
Placeholder: A, for 1. runge(Nonne):
164
Innerooo
124
(Cor MCoer
[Conv2 output)
AtTorme
1.0 range(Nune)::
Foature Voctor
Tradiional Lesning
Computation Stateneent
(Oaline Dauset)
for 1.0 (ange) for k. in Nonec:
cnn
Sel Miterboo
Computation Statement
Mtal--acal
dinemutonn
Hierarchical features Conv2D with full tensor program representation in the search space.
Figure 2:
The performance model's architecture in- cludes two attention blocks that extract coarse and fine- grained features of tensor program, as well as a lightweight MLP layer for directly predicting the score.
Figure 3:
EvolutionSearch() takes the high-level genera- tion sketch GS| and the low-level annotation sam- ples GS2 as input and returns a group of tensor candidates for the cost model training. Next, an evolutionary search strategy is used along with learned cost model to fine-tune the performance of the generated tensor programs. By iteratively mu- generate
neighborhood: and (iii) a simple regression layer for making the final prediction. ATFormer can be easily integrated into existing search algorithms and consistently improve the efficiency of auto- tuning. We believe that the simplicity of our method will attract more research attention to the field of tensor operator optimization, further hancing training and inference efficiency. The fea- ture processing of computation and regression ATFormer is illustrated in Figure The kernel em- bedding layer is composed of two fully connected layers with ReLU activation. The function of the kernel embedding layer is to project the features from low dimension space to new embedding space for similarity measurement. Starting from the batched tensor programs I RLXDin repre- senting a specific type of operator a, where is the accumulated number of the feature statements within I. kernel embedding layer then gener- ates a set of feature statements E E Lx in embedding space. Typically, we use D = 512. The value L is determined by the parameters of high-level structures Φ1 and the low-level details sampling d2 for each subgraph S.
tating high-quality tensor programs, new programs with potentially higher quality. Af- ter a number of measurement trials, the best tensor program configurations can be identified.
Hierarchical Feature Generation. The input of ATFormer is series of mix mix-grained feature vec- tors extracted from Po. where Po is the full tensor program to implement operator a. Each vector represents single computation statement within Ps. These mix- mix-grained fcature vectors are com- posed of two important components: (i) Coarse- Grained operator embedding features that capture the high-level structure of the operator and (ii) Fine-Grained statement features that capture the low-level details of each statement within program Pa. Each operator in the subgraph S can be classi- fied into a few categories, and we represent each operator with a one-hot embedding feature vector that covers all possible operator types. In practice, we use feature vectors of length 10 for the oper- ator embedding and length 164 for the statement features, consistent with the approach used in An- sor (Zheng et al., 2020). The prediction score for subgraph is computed the sum of the predic- tion scores for each innermost loop statement within the loop nests of the full tensor program. More details can found in Figure 2. Model Architecture Our proposed ATFormer model consists of three layers: (i) a kernel em- bedding layer, which extracts compact feature representation; (ii) a computation processing layer, which captures essential information from the in- nermost non-loop computation statements in the
As for the computation layer, statements & Bb. Lout should be split into M stacks of feature statements Z E M firstly. Each stack contains N feature statements of innermost non-loop computation within full tensor program p. We adopt the self-attention mech- anism for feature statements aggregation. With the parameter tensors written Wº. WK W full tensor program with set of innermost non- loop feature statements Z is first encoded into query Q, key K, and value V by three identical linear transformations: Q, K, V = Z W. Then will further calculated by the self-attention
of feature

---

Source
Transfor Learning
Target
Correlation
oGBo
KIX Z0MTEI GPU
Cieepetation Strrmoo
Corpuation
LSTM
Transfrrble
Target Domain 2
Corpuraiio Staemeer
Tesla GPU
Score
ATFormes
Sourre Domain
Targel
Figure 4; Transfer learning among different platforms.
Figure 5: Self-attention between statement features.
layer as:
3.3 Transfer Learning
Q dR
K
Attention(Q2. ( C.K. V)- Softmax
V.
The trade-off between search time and performance improvement is interesting to explore and exploit, as long search times may not always be acceptable. Our current focus is on developing a cost model for optimizing tensor operators ware platform. However, in practical settings, we require cost model that can be used across various hardware platforms. This would allow us to reuse a single-cost model for multiple platforms by provid- ing with new online data during auto-tuning. To achieve this, we pre-train the cost model with an of- fline static dataset and exploit transferable features that are invariant to both source and target domains to speed up the optimization process, depicted in Figure 4. The of transferable features greatly contributes to the success of transfer learning, as different designs may have varying degrees variance. By training the cost model offline using a dataset, we can significantly reduce the frequency of on-device measurements and use the pre- pre-trained parameters as starting point for new search tasks via transfer learning.
(3)
The final prediction of these M tensor programs is computed by regression layer with dimension from 512 1. The predicted score is y E Mx Loss Function The model ranks the performance of potential candidates in large search space. Therefore, the model can trained with ranking losses or regression losses to predict relative or absolute scores. To explore the loss function to train ATFormer, common choice is to use the squared error function regressor which can mostly care about identifying the well-performing tensor programs. The loss function of the model f a full tensor program p with throughput is MSELoss(J, p. h) (E fs) 2) where (p) the set of innermost non-loop computation statements in tensor program p. We train ATFormer as the performance model f. However, we only care about the relative order of tensor program run- time rather than their absolute values during the the following Ran-
specific hard-
compilation. We instead kLoss (Cao et al., 2007) to rank the performance of candidates in the large design space. This can fully exploit the optimal candidates to reduce the impact the search algorithm on final prediction results. The loss function defined as follows:
4 Experiments
4.1 End-to-End Execution Evaluations
Workloads. We evaluate the performance of AT- Former various DNNs, including small and large-scale models. For small-scale models, we AlexNet, VGG-16, MobileNet-V2, R8sNct-18/50 and Bert-Tiny to evaluate the design. As for the large-scale models, we use BERT and GPT-3 mod- els, specifically BERT BERT large- GPT-2targe and GPT-3 We report the the end-to-end in- ference latency with batch size 1 on RTX 2080Ti. Baselines and Settings. For statistic model, we use XGBoost as baseline which has proven to be a state-of-ther feature-basee model in auto-tunin auto-tuning framework (Zheng et al., 2020). For DNN-based
RankLoss
Σ (()
log(1 + e(
(4)
f(i.j) =-sign - u))((((s)) f(s;))
(5)
We the prediction f(x) to select the top- performing implementations of full tensor pro- gram p. The computation graph is trained for tensor programs extracted from all subgraphs. The throughput of all tensor programs is normalized to be in the range of [0, 1].

---

XGRoest TahNei
Inference Intency (mnS)
XGBoost
Nornaiiaal
LSTM
MHA
0.70
ATFormer-I1
ATFormer
AlesNett VGG-16
Mobie-V2
Re-Na-18
ResNer-00
Beri-Tiy
0.60
ATTormer-M
51100
200
300
500
Figure 6: End-to-end performance comparison of cost models across DNNs and normalized by the XGBoost.
dataset size (task #)
Figure 7: Transfer learning performance on TenSet.
learning, we use LSTM with eight heads and 1024 hidden dimensions, and TabNet is implemented in TenSet as another baseline. Note that the search algorithm uses the default configurations, and the search terminates when it runs out of allowed mea- surement trials. We keep the rest of the factors the same for a fair comparison.
on 3090 GPU. The results show that self-attention based models perform best in the final performance compared with the tree-based and DNN-based cost models on different types of GPUs.
Transfer Learning. In Table 1, experiment re- sults on RTX 2080Ti and 3090 show that the pre- trained parameters make the search convergence much faster. With the increasing number of training tasks in the offline dataset from 50 to 500, the learn- ing ability of cost models with self-attention blocks, including MHA, ATFormer-IL. and ATFormer- Mask, become more stable, and they can adapt to the new tasks via transfer learning. ATFormer- series model performs better than the statistic and DNN-based model XGBoost, LSTM in optimized total latency with the parameters trained from TenSet-100 100 TenSer-500. All large-scale mod- els are exported from Hugging Face, with a batch size of and a maximum input sequence length of 512. As shown in Table 2, ATFormer achieves latency speedups of 1.39x, 1.11x, 1.10x, and 1.16x on the 3090 GPU compared to PyTorch runtime. In terms of end-to-end tuning time, AT- Former achieves speedups of 4.97 5.10x,5.0 5.69x. and 6.08x compared traditional learning.
Main Results. Figure 6 shows the final optimized total latency results on the RTX2080Ti GPU. Over- all, the ATFormer-series model performs the best in all cases. Compared with the tree-based model XGBoost, ATFormer outperforms them in all cases with 1.15 1.61 x speedup. Compared with the DNN-based model TabNet, ATFormer outperforms them all cases with 1.14 - 2.14x speedup. Com- pared with LSTM, ATFormer performs equally the best and achieves 0.96- - 1.48x speedup. Although LSTM surpasses ATFormer little in finding the best configuration Bert-Tiny and 16,-16, the amount of computation that can be parallelized in ATFormer leads to shorter used time. Overall, the experiment results from the GeoMean verify the effectiveness of the attention-based modules over the tree- and based performance models.
4.2 Transfer Learning Evaluations
As mentioned in Section 3.3, and 3090 GPUs as different platforms verify our design by two typical metries: i) Fix the measure- ment trails and compare the total latency and ii) Fix a converged latency, and then compare the search time to reach it. To explore transferable features and fast adaptation of auto-tuning between different hardware platforms, AT'Frrmrr is pre-trained with number of samples from TenSet and then fine- tuned using online datasets different platforms. Therefore, we divide our experiment settings into "traditional learning" and "transfer learning" parts.
RTX 2080Ti
The performance of our efficient transfer learn- ing on NVIDIA RTX 3090 GPU can be found in Figure 7. As for the TenSet-50 datasets, curves start from different points at the beginning, and we can find that XGBoost performs best. It mcans that the transferable features in the ATFormer- series models are not fully exploited on the limited dataset #5ask#50) during the training. Obviously, the adaptation skills amplify rapidly with the in- creasing number of tasks the offline dataset. From TenSet-100 100 TenSet-500, we can find that seeormer-series models show fast adaptation and generalization ability across hardware platforms and operators compared with XGBoost and LSTM models.
Traditional Learning. In Table ATFormer achieves the best total latency on RTX 2080Ti, and it performs almost equally best with ATFormer- IL about total latency with a fixed measurement trail
In Table 3, we make the traditional learning and

---

cost madel
XGBoost
LightGDM
LSTM
TabNet
MILA
ATFormer-IL
ATPormer
ATPormer-M
(msss)
latency
time
latency
time
latency
time
latency
time
latency
time
latency
time
lutency
time
Iatency
time
ResNet-18-200p 18-20071
1.47
575
1.58
770
1.29
6634
1.52
748
1.32
687
1.25
706
1.04
787
1.25
762
RTX 208011
TenSes-50
0.98
527
1,02
614
583
1.00
602
1.00
Thandor
TenSer-100
0.96
533
0.98
526
1.07
615
0.82
596
0.87
602
1.00
602
0.85
594
0.84
611
TenSet-200
0.99
536
0.86
525
1.07
611
0.88
582
0.83
602
0.82
612
0.82
604
0.82
632
TenSet 300
0.89
538
0.85
526
1.02
622
0.83
583
0.85
600
0.81
609
0.89
612
0.87
607
TenSes-500
0.96
530
0.81
529
1.03
622
0.82
574
0.83
593
0.87
598
0.84
612
0.79
615
RosNet-18-3000
1.07
589
1.11
676
1.24
762
1.64
741
1.11
6.5
0.97
661
1.02
677
3.01
665
0.74
0.85
0.75
581
605
0.79
604
RTX 3130
Thanndr
TenSet-10U
0.71
540
0.73
526
0.83
599
0.67
620
0.65
607
0.68
601
0.66
606
0.69
614
TenSet-200
0.78
534
0.68
526
0.87
582
0.70
589
0.65
612
0.73
599
0.64
596
0.66
611
Te0Ses3300
0.90
536
0.68
531
0.83
616
0.66
585
0.64
617
0.67
595
0.71
0.66
613
fenSes-500
0.72
535
0.67
540
0.85
618
0.09
587
0.67
591
0.68
581
0.67
607
0.63
609
Table I: Transferable adaptation evaluation between different GPU platforms on ResNet- 18.
cost model
XGBBoot
LSTM
MHA
ATFormer- II.
ATFormer
Speed up
performance (ms ss)
latency
time
latency
time
latency
time
latency
time
latency
time
latency
time
Truditional I. Learning
24.51
3028
32,89
3246
19.13
2890
18.77
2996
17,56
2874
BERT boke
1.39X
4.97
Transfer Learning
23.82
654
33.35
880
19.98
602
19.51
648
18.72
578
Traditional Learning
51.63
5016
59.81
5540
53.21
5218
54,32
5312
46.54
5232
BFRT
1.1xx
5.10x
Transfer Learning
$2.49
60.33
1302
55.88
1084
56.58
1192
47,76
1026
GPT-2arr
Traditional Learning
489.12
6240
502.22
6551
467.22
6311
452.56
6380
445.52
6268
1.10x
5
Transfer Learning
491.24
1392
503.52
1594
468.29
1375
454.18
1272
447.31
1102
Traditional Learning
513.61
7789
542.23
8582
479.42
8082
468.59
7982
442.02
7891
GPT-3p00
1.16x
6.08
Transfer Learning
514.42
543.59
2302
480.12
1890
470.52
1920
443.62
1296
Table 2: The performance of large-scale Transformer models on TenSet-500 with transfer learning.
cost model
XGBoost
LSTM
MHA
ATFormer: 11.
ATFommer
AfFormer M
performance (ms s)
latency
time
latency
time
latency
time
latency
time
latency
latency
time
Traditional Learning
1,26
1026
1.02
1487
1.03
1172
1.20
1269
1.02
1382
1.71
1124
RTX 2080Ti
Transfer Learning
1,23
281
1.08
348
0.99
261
1.15
264
0.99
271
0.93
266
Traditional Learning
1,03
1235
0.79
1125
2054
0.94
2018
RTX 3090
Transfer Learning
0.98
287
1,02
270
0.77
261
0.83
269
0.76
267
0.65
264
Table 3: Pre-trained models on TenSet-500 via transfer learning with converged latency on GPU platforms.
R8
MohileNes-V2
Bert-Tine
Methods
(a)
(b)
(c)
(d)
(o)
(f)
(a)
(b)
(c)
(0)
(o)
(f)
(a)
(b)
(c)
(d)
(e)
(f)
mask?
pre-trained?
4
RMSE Loss?
Rank Loss?
4
4
4
V
AutoTVM?
total Intency (mS)
1.42
1.04
1,23
0.81
0.83
1.92
0.53
0.51
0.76
0.39
0.40
1.29
4.18
3.41
3.97
2.32
2,46
5,07
search time (s)
781
787
762
630
611
3274
962
1000
958
617
604
2996
1127
1141
1150
818
816
3826
Table 4:
Total latency and tuning time of different methods, using ResNet-18, MobileNet-V2 and Bert-Tiny net- works for end-to-end evaluation. The relative gains obtain for batch size = with 300 measurement trials.
transfer learning on different hardware platforms for ResNet- 18 have approximate converged la- tency. ATFormer reduces the search time by up to 5.1x while maintaining the same search quality on RTX 2080Ti. This the best speedup compared with 3.6x by XGBoost, 4.2x by LSTM, 4.8x by ATFormer- IL,, and 2.2x MHA, respectively. Un- der the same conditions, ATFormer also performs the best with reducing the search time by up to 7.7x on RTX 3090 compared with 3.4x by XG- Boost, 4.5x by LSTM, 4.2x by ATFormer- IL, re- spectively. Traditional learning with mask-guided training scheme degrades the performance on total latency and search time. However, transfer learning
with a mask-guided training scheme for ATFormer- Mask performs best in most cases. Comprehen- sive experiments show that it is not easy to make ATFormer-Mask have the approximate converged latency on RTX 2080Ti and 3090 compared with traditional learning and transfer learning. It means that ATFormer-Mas with pre-trained parameters has better task generation for tensor programs and achieves better performance during tuning. Trans- fer learning across different types of CPUs can be found in Appendix A.6.
Overall, ATFormer takes full advantage of trans- ferable features learned from the source domain Tesla T4 GPU and transfers the knowledge to the

---

architecture hehd
hicklen dim
lafency
search time
workloads of DNNs. It shows that the goal of a decent cost model to rank the performance of different tensor programs by relative scores in a given search space.
512
256
1.58
647
4
512
1.24
641
MILA
1024
6
768
1.48
658
8
512
1.19
638
Convergence Speed. In Table 4, method (d) is the proposed ATFormer, which adapts the pre-trained parameters to the new task via transfer learning into method (c). Note that ATFormer with the pre- trained parameters minimizes the total latency of all subgraphs in three DNNs much possible and the search time as quickly as possible. The proposed ATFormer improves the total latency by 4.66x speedup and convergence speed by x.55 speedup. Method (f) is the AutoTVM with lamb- daRank loss function. The performance is inferior to the baseline configuration.
ATFomer-LL
4
512
1.25
706
ATFormer
512
1.04
777
ATFormer-B
Table 5: Different architecture design about ATFormer.
different target domains RTX 2080Ti and RTX 3090 to accelerate the convergence speed with a fixed number of measurement trails. Fast conver- gence is desirable for many users of auto-tuning to have better control of the optimization cost and good performance. For instance, deployment en- gineers may want to obtain optimized model soon possible quickly get an upper-bound estimation of total inference latency in real-world production. They can use the cost model like AT- Former with strong generalization as decent pre- trained parameters to accelerate not only the conver- gence speed but also the total execution inference time. Finally, comprehensive experiments with pre- trained parameters on different sizes of the TenSet dataset show that ATFormer-series models enable fast adaptation in not only cross cross-operator but also cross pross-platform scenarios.
Training Schemes. In Table 4, method (c) incor- porates the mask module into method (b) during traditional learning. Method (d) imports the mask module into method (e) during transfer learning, resulting in notable increase convergence speed. It's worth noting that adding a mask scheme during traditional learning is not very helpful and can even cause decrease in the total latency. However, for transfer learning with pre-trained parameters, incor- porating the mask module is crucial for achieving faster convergence speed. The introduced tech- niques do not require expensive training resources in terms of both time and computation power.
R8
Methods
(+)
(b)
(c)
(d)
(e)
(f)
Model Architectures. Table 5 lists ATFormer with various architectures. To achieve high accu- racy while minimizing the model parameters, we find that the self-attention block, which contains four heads with 512 hidden dimensions, performs the best on the total latency and search time. Note that ATFormer does not benefit from deeper coder layers in the Transformer model. Thanks to its simple and efficient architecture, the inference latency of ATFormer consistently lower than that of the DNNs it optimizes. Thus, we set the two en- coder layers as the final decision. Table 6 shows the relationship between the hierarchical-level features and different architectures to affect total latency and search time on ResNet-18.
Hicrarchical features?
XGBooss?
AlFormer?
w/o Transfer total lantency (ms)
1.47
1.63
1.29
1.58
1.04
L18
w/o Transfer scarch time (s)
573
618
604
787
796
Transfer total latency (ms)
0.96
0.98
1.03
1.12
0,84
0.91
w/ Transfer search time (5)
530
599
622
689
612
632
Table 6: Hierarchical features and performance model architecture improvements for end-to-end evaluation.
4.3 Ablation Study
Various designs are evaluated in this section. We report the performance about total latency, search time on ResNet-18 18 and MobileNet- V2 and accu- racy on the static datasets.
Loss Functions. Table 4 shows two different loss functions in experiments. Method (a) is AT- Former with Root Mean Square Error (RMSE) loss function while method (b) is with lambdaRank loss function. Compared with method (a) and method (b), we find that lambdaRank loss al- ways outperforms RMSE in our design for different
A ccuracy. Table 7 presents the pairwise compari- son accuracy of ATFormer and XGBoost various scales of static datasets. The findings indicate that ATFormer outperforms XGBoost, demonstrating the highest measurement accuracy and providing optimal search quality during the tuning. We suc- cessfully conduct the training process on a server

---

Architecture
XGBoost
ATFormer
ATFormer
ATFormer M
Steiner, Steven Johnson, Kayvon Fatahalian. Frédo Durand, et al. 2019. Learning to optimize halide with tree search and random programs. ACM Trans- actions on Graphics (TOG).
TenSer-50
91.31
93.28
TenSet-300
92.24
90.41
93.82
93.33
TenSet-500
93.08
91.98
94.06
93.71
Table 7: Accuracy of the cost models on TenSet.
Byung Hoon Ahn, Prannoy Pilligundla, Amir Yazdan- bakhsh, and Hadi Esmaeilzadeh. 2020. Chameleon: Adaptive code optimization for expedited deep neu- ral network compilation. In International Confer- ence on Learning Representations (ICLR).
Cood Modiel
TenSa-50
TosSe-100
TesSe-2000
TesSos300 300
TerSai-000
AT
549
916
ATFummer
299
384
588
712
951
ATForren-M
324
416
605
749
912
Table 8: The training time of the ATFormer series cost models during the offline optimization.
Sercan O Ank and Tomas Pfister. 2020. Tabnet: Atten- tive interpretable tabular learning. AAAI Conference on Artificial Intelligence (AAAI).
Model
XGBooo
LSIM
TahteNe:
MHA
AlFommer
AlFommer
ResiNer
3:04
3.35
2,88
1.79
2:51
2.39
Sercan Ô. Arik and Tomas Pfister. 2021. Tabnet: At-
Ben-Ting
15.30
tentive interpretable tabular learning. In Thirty-Fifte AAAI Conference Artificial Intelligence, AAAI 2021. pages 6679-6687 AAAI Press.
Table Traditional learning with different cost models for batch size 8 on the NVIDIA RTX 3090 GPU.
Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yun- ming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. 2019. Tiramisu: A polyhedral compiler for expressing fast and portable code. In 2019 IEFF/ACM International Symposium on Code Generation and Optimization.
equipped with an Intel Core 19-12900K CPU, NVIDIA GeForce RTX 3090 GPU, and a 2TB hard disk. Table 8 presents the specific training times (in seconds) of the ATFormer series models on static datasets. Note that our approach is also suitable for scenarios involving large batch sizes. Table 9 lists experimental results using batch size 8 on the NVIDIA 3090 GPU via traditional learning.
Yang Bai, Xufeng Yao. Qi Sun, and Bei Yu. 2021. Autogtco: Graph and tensor co-optimize for im- age recognition with transformers on gpu. In IEEE/ACM International Conference on Computer- Aided Design (ICCAD).
5
Conclusion
Yang Bai, Xufeng Yao, Qi Sun, Wenqian Zhao, Shixin Chen, Zixiao Wang, and Bei Yu. 2023. Gico: Graph image
This paper introduces ATFormer, a novel and ef- fective design for optimizing tensor programs. AT- Former employs hierarchical features with vary- ing levels of granularity to model the end-to-end compilation. Moreover, self-attention blocks are utilizd explore global dependencies of com- plete tensor program for high-quality evaluation. Through transfer learning, ATFormer achieves faster-converged latency and superior transferabil- ity across different hardware platforms, outperform- ing previous of benchmarks.
and tensor co-design for transformer- recognition on tensor cores. IEEE Transactions on Codputer-Aided Design of Integrated Circuits and Systems.
Bello, Zoph. Q. Le, A. Vaswani. and Shlens. 2019. Attention augmented convolutional networks. In IEEE International Conference Computer Vi- sion (ICCV).
Ali Furkan Biten, Ruben Tito, Andres Mafla, Lluis Gomez, Marçal Rusinol, Ernest Valveny, CV Jawa- har, and Dimosthenis Karatzas. 2019. Scene text visual question answering. In IEEE Conference on Computer Vision and Puttern Recognition (CVPR).
Limitations. We plan to do the transfer learning from GPUs to CPUs and explore the potential of combining with post-training quantization or prun- ing efficiently deploy models. Additionally, we will explore more universal and efficient methods for optimizing tensor programs with ATFormer. This includes leveraging hardware features to opti- mize performance on domain-speciin accelerators, such as NVIDIA's Tensor Cores.
Zhe Cao, Tao Qin, Tie-Yan Yan Liu, Ming-Feng Tsai, and Hang Li. 2007. Leanning rank: from pairwise ap- proach listwise approach. In International Con- ference on Machine Learning (ICML).
Tianqi Chen and Carlos Guestrin. 2016a. Xgboost: A scalable tree boosting system. CoRR.
Tianqi Chen and Carlos Guestrin. 2016b. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, San Fran- cisco, CA, USA, August 13-17, 2016, pages 785-794. ACM.
References
Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Mzu-Mao Li, Michaël Gharbi, Benoit

---

Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lian- min Zheng, Eddie Yan. Haichen Shen, Meghan Cowan, Leyuan Wang. Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018a. TVM: An automated end-to-end optimizing compiler for deep learning. In USENIX Symposium on Operating Systems Design and Implementation (OSDI).
Norman P. Jouppi, Cliff Young, Nishant Patil, David A. Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jef- frey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hag- mann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaf- fey, Alck Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross. Amir Salek, Emad Samadiani, Chris Severn, Gre- gory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing. Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Brick Tuttle, Vi- jay Vasudevan, Richard Walter, Walter Wang. Eric Wilcox, and Doe Hyun Yoon. 2017. In-datacenter performance analysis of tensor processing unit. In IEEE/ACM International Symposium on Computer Architecture (ISCA).
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, al. 2018b. [TVM}: An automated end-to-end optimizing com- piler for deep learning. In USENIX Symposium on Operating Systems Design and Implementation
(OSDI).
Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang. Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018c. Learning to optimize tensor programs. In Annual Conference on Neural Information Processing Systems (NeurIPS).
Jack Choquette, M.-J. Edward Lee, Ronny Krashinsky, Vishnu Balan, and Brucek Khailany. 2021. 3.2 the A100 datacenter GPU and ampere architecture. In IEEE International Solid-State Circuits Conference, ISSCC 2021, San Francisco, CA, USA, February 13- 22, 2021, pages 48-50. IEEE.
Samuel J Kaufman, Phitchaya Mangpo Phothilimthana, Yanqi Zhou, Charith Mendis, Sudip Roy, Amit Sabne, and Mike Burrows. 2020. A learned per- formance model for tensor processing units. In Ma- chine Learning and Systems (MLSys).
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutamova. 2019. BERT: Pre-training of deep bidirectional transformers for language under- standing. In Annual Conference of the North Amer- ican Chapter of the Association for Computational Linguistics (NAACL).
Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu. 2017. Lightgbm: highly efficient gradient boost- ing decision tree. In Advances in Neural Informa- tion Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, De- cember 4-9, 2017, Long Beach, CA. USA, pages 3 3146-3154.
Pratik Fegade, Tianqi Chen, Phillip Gibbons, and Todd Mowry. 2021. Cortex: A compiler for recursive deep learning models. In Machine Learning and Sys- tems (MLSys).
Xiaotian Gao, Wei Cui, Lintao Zhang. and Mao Yang. 2021. Opevo: An evolutionary method for tensor operator optimization. In AAAÍ Conference on Arti- ficial Intelligence (AAAI).
Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017. The tensor algehra compiler. Proceedings of the ACM Pro- gramming Languages.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recog- nition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR).
Mingzhen Li, Yi Liu, Xiaoyan Liu, Qingxiao Sun, Xin You, Hailong Yang. Zhongzhi Luan, Lin Gan, Guangwen Yang. and Depei Qian. 2020. The deep learning compiler: A comprehensive survey. IEEE Transactions Parallel and Distributed Systems (TPDS). Guyue Huang, Yang Bai, Liu Liu, Yuke Wang, Bei Yu, Yufei Ding, and Yuan Xie. 2023. Alcop: Automatic load-compute pipelining in deep learning compiler for ai-s. Proceedings of Machine Learning and Systems.
Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural computation, 1735:773
Jared Roesch, Steven Lyubomirsky, Logan Weber, Josh Pollock, Marisa Kirisame, Tianqi Chen, and Zachary Tatlock. 2018. Relay: A New IR for Machine Leam- ing Frameworks. In ACM SIGPLAN Symposium Programming Language Design & Implementation (PLDI).
Zhe Jia, Marco Maggioni, Jeffrey Smith, and Daniele Paolo Scarpazza. 2019. Dissecting the nvidia turing T4 GPU via microbenchmarking. CoRR, 1bs/1993.77486
Amit Sabne. 2020. Xla Compiling machine learning for peak performance.

---

Mark Sandler, Andrew Howard. Menglong Zhu, drey Zhmoginov, and Liang-Chie Chieh Chen. 2018. Mo- bilenetv2: Inverted residuals and linear bottlenecks. IEEE Conference Computer Vision and Pattern Recognition (CVPR).
Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E Gonzalez, lon Stoica, and Ameer Haj Ali. 2021. TenSet: A Large-scale Pro- gram Performance Dataset for Learned Tensor Com- pilers. In Thirty-fifth Conference on Neural Informa- lion Processing Systems Datasets and Benchmarks Track (Round 1).
Lianlei Shan, Minglong Li. Xiaobin Li, Yang Bai, Ke Lv, Bin Luo, Si-Bao Chen, and Weigiang Wang. 2021. Uhrsnet: semantic segmentation network specifically for uutra-hig--resouutie
images. In IEEE International Conference on Pattern Recogni- tion (ICPR).
Benoit Steiner, Chris Cummins, Horace He, and Hugh Leather. 2021. Value learning for throughput opti- mization of deep learning workloads. In Machine Learning and Systems (MLSys).
Qi Sun, Xinyun Zhang, Hao Geng, Yuxuan Zhao, Yang Bai, Haisheng Zheng, and Bei Yu. 2022. Gtuner: tuning dnn computations on gpu via graph attention network. In ACM/EEEE Design Automation Confer- ence (DAC).
Kai Sheng Tai, Richard Socher, and Christopher Manning. 2015. Improved semantic representations from tree-structured long short-term memory works, arXiv preprint.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances Neural Information Pro- cessing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4- 2017, Long Beach, CA, USA, pages 5998-6008.
Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang, Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer. 2019. Fbnet: Hardware-aware efficient convnet design differ- entiable neural architecture search. In IEEE Confer- Computer Vision and Pattern Recognition (CVPR), pages 10734-10742.
Yu Xiang, Alexandre Alahi, and Silvio Savarese. 2015. Learning to track: Online multi-object tracking by decision making. In IEEE International Conference on Computer Vision (ICCV).
Wenqian Zhao, Yang Bai, Qi Sun, Wenbo Li, Haisheng Zheng, Nianjuan Jiang, Jiangbo Lu, Bei Yu, and Martin DF Wong. 2023. A high- pereoraeffrm accel- erator for super-resolution processing on embedded gpu. IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems.
Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu. Cody Hao Yu. Ameer Haj-Ali, Yida Wang, Jun Yang. Danyang Zhuo, Koushik Sen, et al. 2020. Ansor: Generating nigh-performance tensor programs for deep learning. In USENIX Symposium on Operating Systems Design and Implementation (OSDI).

---

A Appendix
Arithmetic intensity curve: We only sam- ple 10 points from curve which is defined as FLOPs which similar to the roof-line Hytee
A.1 Feature Extraction Details
The feature before ATFormer training can be rep- resented as two different granularities: coarse- grained and fine-grained levels. The coarse-grained level feature can describe each search task in the computation graph. It has 10 elements with the one-hot encoding pattern. In our specific code implementation, the carse-grained vector contains these oaer. "min", "add", "ConrtdOutput", "Conv2d = "Depth- wiseConv2d", "dense", "softmax", "compute(b, i, i)" The "max" and "min" can represent some acti- vation functions in deep learning. "dense" means the fully connected layer in computation graph and "compute(b; i, j)" is very important function to implement each tensor operation in deep learning. If the intermediate representation about some op- erators are fused into the same "compute(b, j)) primitive, means these operators are fused to- gether and can run very efficiently on the specific hardware platforms. As for the fine-grat tineggraindd vec- tor, the length of it including all the listed features for statement is 164. We use the same set of features for both Turing 2080Ti and Ampere 3090 GPUs. It can be summarized follows:
model used in computer architecture. It can help us to recognize the type of the search task or operator in computation graph such as compute-intensive or memory-intensive oper- ator on a specifie hardware platform.
We perform fea- ture extraction for at most five buffers. It includes "Access type". "Bytes", "Unique bytes", "Lines", "Unique Lines", "Reuse type", "Reuse distance". "Reuse counter", "Stride", "Accessed bytes divided by reuse".
Buffeer Access Feature:
Allocation related features: The size of the allocated buffer for the output results of each statement in a full tensor program.
With the combination of coarse boarse-grained and fine- grained feature vectors, we can construct them into hierarchical feature vector to take full advantage of each statements in a full tensor program.
A.2 Implementation Details
ATFormer is implemented on the top of Ansor and evaluated from two aspects: end-to-end search ef- ficiency and quality, well as performance porta- bility. We compare ATFormer against the state-of- the-art methods, including both the statistic and DNN-based cost models. The items labeled with XGBoost represent the Ansor default configura- tion. We also provide detailed ablation study of the model architecture, accuracy, loss function, convergence speed, and training scheme, with in- sights and qualitative results. The generated ten- sor programs are evaluated on two different GPU architectures: Turing RTX 2080Ti and Ampere RTX 3090, with float32 data types used for all evaluations. We train the cost model using the Adam optimizer for 50 epochs, with a starting learn- ing rate of -4 that decays to le -6, and a train- ing batch size set to 512. We use TVM v0.8dev in TenSet (Zheng et al., 2021), LLVM 1 11.0, and CUDA 11.0 for compilation, while XGBoost 1.5.0 and PyTorch 1.7.1 are used for training models. The use of "mask" is widely adopted technique for training transformers. In Figure 5, each tensor program transformed into a sequence of vectors. with each vector representing a tensor computa- tion statement. During training, all sequences are of the same length, and any shorter sequences are
Number of float operations: The number of addition, subtraction, division, modulo oper- ation, less-than, greater-than, intrinsic math function such as exp, sqrt.
Number of integer operations: Similar to the number of float operations, but for the operations with inte- ger operations.
Vectorization related features: The number of the innermost vectorized loop statements in a full tensor program.
Unrolling related features: The number of the innermost unrolling loop statements in full tensor program.
Parallelization related feature: The num- ber of the innermost parallelization loop state- ments in full tensor program.
GPU thread binding related features: The lengths of blockldx.x, blockldx.y, blockldx.2, threadldx.x, threadldx.y, threadldx.2 and vir- tual threads which can avoid bank conflict problem in shared memory.

---

padded with zeros the end. The padded items are masked out and excluded from the loss com- putation. Our ablative models, including MHA, ATFormer- II, ATFormer, and ATFormer-M, were also experimented with. MHA the basic Multi- Head Attention layer, ATFormer- IL only has one encoder layer, ATFormer has two encoder layers. and ATFormer-M uses the "mask" scheme during training.
ResWet
Oura tranater learning
XoBoest cranster carning
24-
Deliast
223
1.-
200
600
68.3
006
1200
Figure 8: Convergence analysis on ResNet-18.
A.3 Dataset Details
We evaluated our design using TenSet, a large-scale and challenging dataset for search-based tensor compilers. TenSet comprises 52 million perfor- mance records of tensor programs obtained from real measurements on different hardware platforms. Various randomly generated tensor programs for popular workloads are compiled via the TVM com- piler and executed on the target hardware platforms. To ensure the inclusion of diverse workloads essen- tial for generalization ability, we collected tensor programs from 120 networks with 13,848 tasks on the NVIDIA Tesla T4 GPU. This dataset serves as a series of static offline datasets.
We use float32 as the data type for all evalua- tions. We train our model with the Adam optimizer for 50 epochs with a starting learning rate of 7e the learning rate decays to le and the training batch size is set to 512. We use TVM v0.8dev in TenSet, LLVM 11.0 and CUDA 11.0 for com- pilation. Meanwhile, we use XGBoost 1.5.0 and PyTorch 1.7.1 for training models.
To explore transferable features and fast adap- tation of ATFormer between different hardware platforms, ATFormer is pre-trained using offline learning with number of samples from TenSet, and then fine-tuned using online learning dif- ferent platforms. For the offline learning, we domly sample 50, 100, 200, 300, 500 search tasks from TenSet NVIDIA Tesla T4 GPU.
A.4 Benchmark Details
We evaluate the performance of generated pro- grams by ATFormer on two levels: end-to-end network evaluations and performance portability via transfer learning. For each level of evaluation, we compare ATFormer against the state-of-the-art methods, including the statistic models:
We train 40 models including XGBoost, Light- GBM, LSTM, TabNet, Multi-head attention, ATFormer- IL, ATFormer, ATFormer-Mask for all of experiment evaluation in this paper. Due to the limitation of maximum file size (100MB) in supple- mentary material, we release the pre-trate pre-trained model offline learning by Tenset-500 for AFTormer- IL. ATFormer, ATFormer-Mask Multi-head attention and TabNet. All of the pre-trained models for XG- Boost. And we release running scripts in the supple- mentary material to reproduce the results in Section 5 Table 1. More details about the hyperparameters of each cost model in our experiments can be found in Table 12, Table 13, Table 14, Table 15, Table 16, Table 17, and Table 18.
XGBoost (Chen and Guestrin, 2016b)
LightGBM (Ke et al., 2017)
and DNN-based models:
LSTM (Hochreiter and Schmidhuber, 1997)
Multi- Attention (Vaswani et al., 2017)
TabNet (Arik and Pfister, 2021)
A.5 Convergence Analysis.
The generated tensor programs are benchmarked on two different architecture GPU platforms:
In Figure 8, we present the tuning trials-latency curves that illustrate various stages of auto-tuning with different configurations on ResNet-18. We performed four types of experiments on ResNet ing two settings: with transfer learning and without transfer learning. The blue line indicates ATFormer with transfer learning to expedite the tuning pro-
NVIDIA 2080Ti GPU with Turing architec- ture (Jia al., 2019)
NVIDIA 3090 GPU with Ampere architec- ture (Choquette et al., 2021)

---

cost model
XGBoost
1.TTM
MHA
ATForerr-I II.
ATFormer
performance (ms s)
latency
time
latency
time
lateney
time
latency
time
latency
time
Traditional Learning
6.13
658
6.16
731
6.12
642
6.22
ResNet-18
Transfer Learning
6.16
334
6.25
451
6.19
6.29
6.18
Traditional Learning
19.59
652
21.23
697
17.50
630
17.52
614
16.90
643
ResNet-50
Trunsfer Leurning
20.01
342
21,99
461
18.11
338
17.91
362
17.02
318
Traditional Learning
36.92
891
39.85
1004
35.69
839
34.51
826
30.01
VGG-16
Transfer Leurning
37.51
395
40.17
422
36.79
34.87
Traditional Learning
17.98
1012
19.22
1433
17.55
1126
16.09
1168
15.11
1232
BERT-Tiny
Transfer Learning
18.05
396
19.57
498
17.91
401
16.41
416
15.16
388
Table 10: Pre-trained models with the converged latency on the CPU platforms.
cost model
XGBoost
LSTM
MHA
ATFormer-I II.
ATFommer
pertormance (ms /s)
latency
time
latency
time
latency
time
latency
time
latency
time
Traditional Learning
5.28
634s
5.91
702
5.17
611
5.32
602
4.75
628
ResNet-18
Transfer Learning
5.21
314
5.88
432
5.16
326
5.19
384
4.74
254
Traditional Learning
16.42
621
18.23
632
13.51
608
12.51
584
11.62
602
ResNet-50
Transfer Learning
342
21.99
18.1
338
17.91
362
17.02
318
Traditional Learning
29.52
845
31,54
967
28.55
799
28.71
796
25.49
812
VGG-16
Transfer Learning
29.41
352
31,47
378
28.46
299
28.69
278
25.46
216
Traditional Learning
13.88
862
15.22
1138
13.55
986
14.41
942
11.55
998
BERT-Tiny
Transfer Learning
13.76
339
15.47
438
13.91
345
14.39
377
11.58
Table II:
Pre-trained models with the converged latency the Tensor Cores.
cess. We observe that the converged latency is the best among the four configurations. The orange line represents the same tuning process with the XGBoost cost model, and we note that the con- verged latency is inferior to the one with ATFormer. The green line shows ATFormer without transfer learning, and we can observe that the convergence speed is exceptionally fast. The red line represents the Ansor optimization, and observe that the convergence speed and the final converged latency are both less than the ones achieved by the green line with ATFormer. Therefore, we can infer that ATFormer can expedite the tuning process com- pared to traditional learning methods through trans- fer learning and outperforms the state-of-the-art tensor compiler Ansor.
dimension scalar.
Model Name
Parameter
Value
max_ max__epth
6
gamma 0.003
XGBoost
sii_chill__eiie
2
0.2
Table 12: Hyperparameters of XGBoost.
Model Name
Parameter
Value
num leavle
72
boosting
gbdt
Ir bagging_free
0.16
LightGBM
The main components in ATFormer model archi- tecture can be categorized into three layers:
min sum Icaf
fraction_f
0.84
fraction
0.89
Kernel embedding layer: The function of kernel embedding layer is to change the 164 10 dimensions into 512 dimensions.
Table 13: Hyperparameters of LightGBM.
The function of compu- tation layer is to obtain the relationship be- tween each innermost non-loop statement in loop nests of full tensor program.
Computation layer:
A.6 Other Platforms: Intel CPUs
We use the dataset from Intel Platinum-8277 to verify transferability on Intel E5-2698 CPU with fixed converged latency (6.13ms) by the same measurement trials for ResNet-18. More details can be found in Table 10. Therefore, ATFormer also works well for CPU with lots of different DNN
Regression layer: The function of regression layer to project the final prediction about each innermost non-loop statement an one

---

benchmarks including ResNet-50, VGG-16, BERT- Tiny with batch size 1. As for the ResNet-18, we fix the converged latency to 19.59ms, the traditional learning will cost 658s to search the optimal config- uration with XGBoost performance model. But the ATFormer can search the optimal implementation of ResNett-50 with 643s by the same measurement trials under the 16.90ms converged latency. We can get the conclusions from the VGG-16 and BERT-Tiny neural networks.
Model Name
Parameter
Value
in dim
164 + 10
hidden
256
out_dim dim
1
n_d
TabNet
n_a
n_ steps
gamma
7e
Table 14: Hyperparameters of TabNet.
A.7 Performance Tensor Cores
The recent advancements of GPU hardware tech- nology have resulted a significant increase computing power, particularly with the introduc- tion of the Tensor Cores on NVIDIA GPUs. Unlike the scalar-to-scalar primitives found in CPUs the general CUDA Cores in GPUs, Tensor Cores provide specialized tensor computation capacities, which deliver over 10x higher throughput. No- tably, the initial version of Tensor Core designed for handling the GEMM with half-precision input and full preciison output. Recently, new features supporting different datatypes such as int8, int4 and int1 input variables have been introduced in the latest architecture (Truing and Ampere), The collection process takes 5 days with server equipped with an Intel Core 19-1290 CPU and NVIDIA GeForce RTX 3090 GPU. The sampling selection process for the operator conducted in a manner similar to that on the GPU's CUDA cores. We use the floating point 16 (fp16) as the experiemnt datatype and additional experimental results on transfer learning are presented in Ta- ble 11.
Model Name Parameter
Value 164 + 10 1024
in_ddni dim
hidden_din
LSTM
out_dim
Table 15: Hyperparameters of LSTM.
Model Name
Parameter
Value
in rn_dim num heads
164 10
MHA
hidden_ dim
512
out _im
Table 16: Hyperparameters of Multi-head Attention.
Model Name
Parameter
Value
in in_dim num heads hidden dim num layers out out__im dim padding mask
164 + 10
4
512
ATFormer
2
False
Table Hyperparameters ATFormer.
Model Name
Parameter
Value
in_dim num headh hidden_diin
164 + 10
512
ATFormer-M
num layers
2
out_ddi padding mask
True
Table 18: Hyperparameters of ATFormer.


Reviews:
[{"review_id": "Sdpnk15khM", "reviewer": "Reviewer_astq", "summary": "", "questions": "A) What is method (e) in line 549? \n\nB) Why might performance for LSTM be better on Bert-Tiny and VGG-16? That could be an interesting insight. \n\nC) The batch size is notably set to 1. Are the insights of this paper applicable to larger batch sizes? \n\nD) Are there experimental details/settings for the online datasets? They appear to be dynamic with execution, but it may be helpful to include details of their size/settings for reproducibility. \n", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your comments as well as the appreciation of our work. Our responses to your concerns are demonstrated below:\n\n**Q1:** What is method(e) in line 549?\n\n**A1:** Method (e) corresponds to the utilization of ATFormer with two attention blocks (ATFormer-2L) that are trained through transfer learning and rank loss to fine-tune the tensor program. The distinction between method (e) and method (d) lies in the fact that method (d) incorporates a masking mechanism known as ATFormer-M during the training phase.\n\n**Q2:** Why might performance for LSTM be better on Bert-Tiny and VGG-16?\n\n**A2:** Our LSTM model is relatively simple, consisting of only one layer with a hidden dimension of 1024. As a result, it has fewer parameters compared to the ATFormer-series models. In contrast to the variable kernel sizes used in models like ResNet or MobileNet, the VGG-16 network has fixed kernel sizes in its convolutional layers, making its architecture straightforward. This stability in the predictions of LSTM differentiates it from the interweaving of different convolutional layer sizes. Similarly, in the Bert-Tiny model, the dimensions of the optimization matrix for Q, K, and V in the encoder remain unchanged. This characteristic allows LSTM to achieve slightly better performance in terms of cost model updates using online datasets compared to the ATFormer series. Moreover, another potential factor is that the construction of the online dataset is relatively small in scale. When training a cost model from scratch, the presence of fewer parameters in the cost model can facilitate quicker convergence of its performance. However, it should be noted that these improvements are mainly observed in traditional learning with online datasets. In transfer learning experiments, particularly in optimizing NLP-based models, LSTM cost model does not outperform the superior performance achieved by ATFormer cost models, regardless of the total tuning time or final optimized inference latency.\n\n**Q3:** The batch size is notably set to 1. Are the insights of this paper applicable to larger batch sizes?\n\n**A3:** Thanks for pointing out this interesting part. We would like to highlight that our approach is also suitable for scenarios involving large batch sizes. In order to validate our design within the given time constraints during the rebuttal phase, we perform experiments using 8 batches on the NVIDIA 3090 GPU via traditional learning. The experiments are conducted on Bert-Tiny and ResNet-18, utilizing the same baselines and settings as outlined in our original paper.. More details can be found in the following tables.\n| batch = 8   | ResNet-18 (ms) | Bert-Tiny (ms) |\n| ----------- | -------------- | -------------- |\n| XGBoost     | 3.04           | 17.42          |\n| LSTM        | 3.35           | 14.83          |\n| TableNet    | 2.88           | 16.98          |\n| MHA         | 2.79           | 15.32          |\n| ATFormer-1L | 2.51           | 15.49          |\n| ATFormer-2L | **2.39**       | **14.37**      |\n\nIn the revised version, we intend to expand our analysis by including additional batch sizes such as 2, 4, and 16 for both transfer learning and traditional learning scenarios. This will provide a more comprehensive understanding of our design across a range of batch sizes.\n\n**Q4:** Are there experimental details/settings for the online datasets? They appear to be dynamic with execution, but it may be helpful to include details of their size/settings for reproducibility.\n\n**A4:** Thanks for the comments. ATFormer is built on top of Ansor[1], which employs an online data collection approach during the search process to construct the datasets. In order to maintain a fair comparison with the baselines, we have continued to employ the same approach as Ansor for constructing online datasets. Our approach also combines an evolutionary search algorithm with iterative construction of online datasets. In each iteration, we utilize newly resampled programs and good programs from previous iterations as the initial population for the evolutionary search. We use a set of evolution operations, including *Tile size mutation*, *Parallel mutation*, *Pragma mutation*, *computation location mutation*, and *node-based crossover*, to rewrite and generate tensor programs. The evolutionary search leverages mutation and crossover to repeatedly generate new candidate programs over multiple rounds, ultimately producing a small set of tensor programs with the highest scores. These selected tensor programs are then compiled and executed on a GPU platform to measure the real inference latency. The collected measurement data is subsequently utilized to update the cost model, gradually enhancing its accuracy to better match the GPU performance with the online datasets. For more detailed information, please refer to the hyperparameters specifically designed for the evolutionary search in collecting the online datasets:\n| Parameters                        | Value    |\n| ---------------------- | -------- |\n| sample_init_min_population        | 50       |\n| sample_init_use_measured_ratio    | 0.2      |\n| evolutionary_search_population    | 2048     |\n| evolutionary_search_num_iters     | 4        |\n| evolutionary_search_mutation_prob | 0.85     |\n| gpu_multi_level_tiling_structure  | SSSRRSRS |\n| max_innermost_split_factor        | 64       |\n\nThe term \"gpu_multi_level_tiling_structure\" denotes the utilization of the \"SSSRRSRS\" template to generate the high-level structure of tensor programs. In this template, \"S\" represents one tiling level of space loops, while \"R\" represents one tiling level of reduction loops. The loops in the first three space tilings are bounded to BlockIdx, virtual thread (to reduce bank conflicts), and ThreadIdx on the GPU, respectively.\n\nTo facilitate reproducibility and further development on ATFormer, we will open source all of code including training and testing, as well as pretrained models.\n\n[1] Zheng, Lianmin and Jia, Chengfan and Sun, Minmin and Wu, Zhao and Yu, Cody Hao and Haj-Ali, Ameer and Wang, Yida and Yang, Jun and Zhuo, Danyang and Sen, Koushik and others, 14th USENIX symposium on operating systems design and implementation (OSDI 20) \n\n**Q5:** Typos and figures/tables size comments.\n\n**A5:** Thanks for the comments. We will fix all the typos in the revised version. In the original submission, certain figures and tables were compressed to comply with the page limitation. In the revised version, we will carefully condense some of the content and move it to the appendix, thus creating more appropriate space for these tables and figures.\n\n**Q6:** It would be helpful to define ATFormer-1L/M earlier in paper and clearly outline of (a) - (f).\n\n**A6:**  We totally agree with you and we provide more elaborate descriptions below. The model we refer to as ATFormer-1L is composed of a single layer of attention block, distinguishing it from the standard ATFormer cost model, which encompasses two layers of attention blocks. Furthermore, ATFormer-M indicates the utilization of a mask mechanism during training to fortify the stability and performance of ATFormer. Method (a) refers to training ATFormer using rmse loss via traditional learning to optimize the network's inference performance. Method (b) involves substituting the rmse loss in method (a) with rank loss. Method (c) builds upon method (b) by incorporating ATFormer-M, which integrates the mask mechanism, to further enhance the performance of cost model. Method (d) extends method (c) by integrating transfer learning with pretrained-models to expedite the tuning process of the network. Method (e) combines the standard ATFormer via transfer learning without the mask mechanism for optimization. Method (f) represents the baseline approach using AutoTVM for operator tuning. In this method, only rank loss is employed, and the operator is optimized through traditional learning. Finally, in the revised version, we will integrate these variants into the appropriate sections to facilitate readers' comprehension of our design and configuration.\n\n**Q7:**  While 2 GPUs were tested, evaluating on more GPUs with variance in architecture would strengthen the claims.\n\n**A7:** Thanks for the comments. In the future work, we will extend and develop the ATFormer to a broader range of architectures, including NVIDIA Jetson AGX, Xavier, NX, TX2, as well as the emerging hopper architecture, to accelerate a wider variety of operators and models.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The research problem of interest (neural network compilation for efficiency) has critical implications for practitioners. The authors evaluate a wide range of DNNs (MobileNet/ResNet/BERT/GPT to name a few) and make a solid, unique effort in this direction with their approach.\n\nThe approach is well-motivated based on challenges with past approaches (large search space, cross-device transferability, long optimization times). \n\nThere are adequate baselines (XGBoost, LSTM, TabNet) that the approach outperforms on many models, and the ablations are insightful.\n\nThe testing across other GPU platforms (NVIDIA 2080 Ti with Turing / NVIDIA 3090 GPU with Ampere) is interesting and practical. \n", "reasons_to_reject": "Some of the gains seem to be incremental (Fig. 6); LSTM outperforms ATF on 2 out of the 6 models in inference latency. \n\nOne of the primary claims of the paper is that the approach demonstrates effective transfer across devices. While 2 GPUs were tested, evaluating on more GPUs with variance in architecture would strengthen the claims. \n", "questions_to_authors": "A) What is method (e) in line 549? \n\nB) Why might performance for LSTM be better on Bert-Tiny and VGG-16? That could be an interesting insight. \n\nC) The batch size is notably set to 1. Are the insights of this paper applicable to larger batch sizes? \n\nD) Are there experimental details/settings for the online datasets? They appear to be dynamic with execution, but it may be helpful to include details of their size/settings for reproducibility. \n", "paper_topic_and_main_contributions": "In order to optimize deep neural network efficiency and overcome challenges such as long optimization times and deployment transfer to other hardware, this paper proposes a cost modeling and tensor program optimization framework, ATFormer. \n\nThe tuning problem is framed as an optimization problem in a hierarchical search space. ATFormer uses attention modules to capture global/long-range dependencies in features. \n\nThe model is evaluated in terms of traditional (online) and transfer learning, with results indicating improved latencies and search times across GPU devices.\n\nOther contributions include extensive performance analysis of auto-tuning with different strategies (tree-based/DNN-based approaches).\n", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "There are a fair amount of typos, which warrants further proofreading of the entire paper. Noting some:\n- Line 53 “an” should be “can”\n- Line 231 “perdiction” \n- Line 528 “AFTormer”\n- Figure 7 X-axis overlaps \n\nIn general, the figures/table text should be larger.\n\nIt would help to define ATFormer-1L and ATFormer-M/Mask earlier in paper. \n\nTable 4 and 6 would benefit from a clearer outline of (a) - (f) in the text somewhere. \n"}}, {"review_id": "LmGgVd2pkB", "reviewer": "Reviewer_y1Wr", "summary": "", "questions": "1) What are the high-level structure and low-level structure in hierarchical feature generation?\n\n2) In Figure-5, how to incorporate statement features into self-attention and how to update self-attention after scanning the full tensor program?\n\n3) Why is the throughput metric used in the loss function although the experiment results show a latency improvement? \n\n4) Table-4 should include one or two nlp models.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and appreciation of our work. To relieve some of your concerns, we have posted our responses below:\n\n**A1:** Thanks for the comments. Our method does not solely optimize the innermost statements but rather focuses on optimizing the full tensor program, including conventional loop-tiling, reorder,unrolling and so on. In Section 3.1 of the problem formulation, we define a hierarchical search space consisting of two important components: high-level structure and low-level details. The high-level structure is responsible for making structural changes to the entire for-loop nest, including loop tiling, reordering, unrolling and so on. On the GPU, we utilize the \"SSSRRSRS\" template to generate the corresponding high-level structure. Here, \"S\" refers to one tiling level of space loops, and \"R\" represents one tiling level of reduction loops. The loops in the first three space tilings are bound to BlockIdx, virtual thread (for reducing bank conflicts), and ThreadIdx on the GPU, respectively. Different transformations in the outermost for-loops will result in different features extracted from the innermost assignment statement, leading to varying throughput predictions and learning processes in the cost model. Overall, our method considers various optimizations of the outermost loops during the optimization process. By extracting features from the innermost non-loop statements, we train a learnable cost model on the GPU to automatically optimize the full tensor program. We train the cost model and mainly target the data parallel tensor programs which consist outermost loop-nest and some innermost assignment statements. We use the ATFormer to predict the score of each non-loop innermost statement. An example are shown as follows:\n```python\nfor i in range(10):\n  for j in range(10):\n      B[i][j] = A[i] * 2                          # Statement B\nfor i in range(10):\n  \tC[i] = B[i][i] - 3                          # Statement C\n```\nIn this example, there are two non-loop innermost statement, B and C. We estimate the cost of these tensor programs. The total cost of it is the sum of cost B and cost C. We perform static analysis to extract the features for every non-loop innermost statement and make prediction for them. These predictions, along with features extracted from the entire program, such as the number of *used cached lines*, *used memory*, *reuse distance*, and *arithmetic intensity curves*  are used to estimate the cost of the entire program. In the following code snippet, we utilize an example of matrix multiplication and ReLU function to elucidate the high-level structure and low-level details as mentioned before.\n```python\nfor i in range(512):\n  for j in range(512):\n    for k in range(1024):\n      C[i, j] += A[i, k] * B[k, j]\nfor i in range(512):\n  for j in range(512):\n    D[i, j] = max(C[i, j], 0)\n```\nOne of the possible high-level structure transformation is shown as below:\n```python\nfor i.0 in range(Tile_I0):\n  for j.0 in range(Tile_J0):\n    for i.1 in range(Tile_I1):\n      for j.1 in range(Tile_J1):\n        for k.0 in range(Tile_K0):\n          for i.2 in range(Tile_I2):\n            for j.2 in range(Tile_J2):\n              for k.1 in range(Tile_I1):\n                for i.3 in range(Tile_I3):\n                  for j.3 in range(Tile_J3):\n                    C[...] += A[...] * B[...]\n        for i.4 in range(Tile_I2 * Tile_I3):\n          for j.4 in range(Tile_J2 * Tile_J3):\n            D[...] += max(C[...], 0.0)\n```\nThe following two code snippets demonstrate the low-level details that can be generated based on the aforementioned high-level structure in the hierarchical search space.\n```python\nparallel i.0@j.0@i.1@j.1 in range(256):\n  for k.0 in range(32):\n    for i.2 in range(16):\n      unroll k.1 in range(16):\n        unroll i.3 in range(4):\n          vectorize j.3 in range(16):\n            C[...] += A[...] * B[...]\n  for i.4 in range(64):\n    vectorize j.4 in range(16):\n      D[...] = max(C[...], 0.0)\n```\n```python\nparallel i.2 in range(16):\n  for j.2 in range(128):\n    for k.1 in range(512):\n      for i.3 in range(32):\n        vectorize j.3 in range(4):\n          C[...] += A[...] * B[...]\nparallel i.4 in range(512):\n  for j.4 in range(512):\n    D[...] = max(C[...], 0.0)\n```\nBy examining the provided code snippets, it becomes evident that throughout the entire process of training and prediction in the cost model, there are considerations for common optimizations such as tiling that affect the outermost for-loops. To provide a more comprehensive understanding of our design and optimization methodology, we will showcase detailed code examples in the appendix. These examples will illustrate the process and incorporate common optimizations to enhance readers' comprehension.\n\n**Q2:** What are the high-level structure and low-level structure in hierarchical feature generation?\n\n**A2:**  Each tuned operator in the neural network can be represented as a loop-nested structure in their implementation and then mapped to the specific hardware accelerators for the efficient inference or training. As for the high-level structure, it represents the coarse-grained statement features and it is labeled as a tag such as dense, softmax or conv2d to describe the type of this operator as the graph-level representation. As for the low-level structure, it represents more details about the implemenation and optimization of the operator such as multi-level tiling, reordering and unrolling as the tensor-level representation. It is worth noting that the concepts of high-level and low-level in the hierarchical search space defined in Section 3.1 of the problem formulation differ from the concepts outlined in Section 3.2  hierarchical feature generation. The former emphasizes the overall transformations of tensor program, while the latter pertains to the representation and extraction of features within the innermost assignment statements. We appreciate the reviewer's valuable feedback and inquiries regarding this matter. In order to facilitate a deeper understanding of our work, we will further emphasize these two aspects in the revised version, ensuring that readers can grasp them more effectively.\n\n**Q3:** In Figure-5, how to incorporate statement features into self-attention and how to update self-attention after scanning the full tensor program?\n\n**A3:** Figure 5 portrays a high-level diagram that showcases the utilization of self-attention for relationship modeling, where the extracted features from the innermost assignment statement are fed as input.  Figure 3 offers a more lucid depiction of the features extracted from the statement and their subsequent incorporation into self-attention. Initially, a feature vector of dimension 164 will be extracted  for the innermost assignment statement including computation, memory access, and arithmetic strength. Subsequently, this vector will undergo concatenation with a 10-dimensional one-hot operator category feature, resulting in the creation of a new feature vector with a length of 174. Next, the embedding layer will embed each innermost assignment statement extracted from different full tensor programs into a new dimensional space representation. Finally, these embeddings will be multiplied by learnable weight matrices to obtain the corresponding Q, K, and V matrices. This comprehensive process enables end-to-end training and facilitates the updating of the self-attention mechanism.\n\n**Q4:** Why is the throughput metric used in the loss function although the experiment results show a latency improvement?\n\n**A4:** Thanks for the comment. The cost model is used to predict the fitness of each tensor program, which is the throughput of one tensor program. During the process of model inference on a single machine with a single GPU, it is generally assumed that there exists an inverse relationship: ${throughput} ∝ \\frac{1}{latency}$. It means that as the throughput increases, the inference latency decreases. Therefore, there is a prevailing benchmark standard for performance evaluation of the entire end-to-end network lies in the inference latency in both Academia and industry when developing a new inference engine or framework. Certainly, if throughput is employed exclusively during experimental evaluations, it will yield identical conclusions to those obtained using inference latency as the metric.\n\n**Q5:** Table-4 should include one or two nlp models.\n\n**A5:** Thanks for the comments. Due to the limited rebuttal time, we add the nlp-based model Bert-Tiny in the table below to further verify our design. \n|      Bert-Tiny       |          |          |          |          |          |          |\n| :------------------: | -------- | -------- | -------- | -------- | -------- | -------- |\n|       Methods        | (a)      | (b)      | (c)      | (d)      | (e)      | (f)      |\n|        Mask?         |          |          | &#10004; | &#10004; |          |          |\n|     Pre-trained?     |          |          |          | &#10004; | &#10004; |          |\n|      RMSE Loss?      | &#10004; |          |          |          |          |          |\n|      Rank Loss?      |          | &#10004; | &#10004; | &#10004; | &#10004; | &#10004; |\n|       AutoTVM?       |          |          |          |          |          | &#10004; |\n| Total Latency （ms） | 4.18     | 3.41     | 3.97     | **2.32** | 2.46     | 5.07     |\n|   Search Time (s)    | 1127     | 1141     | 1150     | **818**  | 816      | 3826     |\n\nIn the revised version, we will integrate a wider nlp-based models to facilitate the comprehension of our design for the readers.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1) The work is well-written, and the sufficient experiment results demonstrate the soundness of the authors' design.", "reasons_to_reject": "1) The problem chosen for this study is simple, with only the innermost nonloop statement to be tuned.", "questions_to_authors": "1) What are the high-level structure and low-level structure in hierarchical feature generation?\n\n2) In Figure-5, how to incorporate statement features into self-attention and how to update self-attention after scanning the full tensor program?\n\n3) Why is the throughput metric used in the loss function although the experiment results show a latency improvement? \n\n4) Table-4 should include one or two nlp models.", "paper_topic_and_main_contributions": "The training and inference efficiency of deep neural networks heavily relies on the performance of tensor operators on specific hardware platforms. Compilation-based optimization flows, including automatic tensor generation and parameter tuning, are necessary for efficient model deployment. 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. ATFormer enables rapid adaptation of performance tuning across various GPU platforms using pre-trained parameters on static datasets.\n\nCompared with the state-of-the-art baselines, ATFormer can predict the optimal implementation of tensor operators to reduce latency time with minimal effort on modern DNN benchmarks.", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "umTFqZNI1S", "reviewer": "Reviewer_fHck", "summary": "", "questions": "1.\tIn Figure 3, why are the outputs of self-attention and embedding concatenated again? What is the reason for including an embedding layer in every attention block? Additionally, could you confirm whether the attention block utilizes pre- and post-layer norms?\n2.\tPlease provide the latency values for BERT-large and GPT models in Figure 6? It seems that ATF offers little improvement compared to MHA in some cases.\n3.\tIt seems that ATFormer with 2 attention blocks performs worse than ATFormer-1L on the dataset with a size of 300 points in Figure-7. Can you provide insights into the reasons behind this performance difference?\n4.\tIn Table-2, it appears that ATF reduces search time by utilizing transfer learning, but this comes at the cost of increased kernel execution latency. Does this ultimately lead to a worse overall result?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your valuable comments on our work. To address your concerns, our responses are demonstrated below:\n\n**Q1:** The code generation templates are overly simple. The final prediction/tuning objective only considers the innermost non-loop statements. There is no optimization like loop tiling in the code generation task, which is shown to be important for code optimization of DNN workloads.\n\n**A1:** Thanks for the comments. Our method does not solely optimize the innermost statements but rather focuses on optimizing the full tensor program, including conventional loop-tiling, reorder,unrolling and so on. In Section 3.1 of the problem formulation, we define a hierarchical search space consisting of two important components: high-level structure and low-level details. The high-level structure is responsible for making structural changes to the entire for-loop nest, including loop tiling, reordering, and unrolling. On the GPU, we utilize the \"SSSRRSRS\" template to generate the corresponding high-level structure. Here, \"S\" refers to one tiling level of space loops, and \"R\" represents one tiling level of reduction loops. The loops in the first three space tilings are bound to BlockIdx, virtual thread (for reducing bank conflicts), and ThreadIdx on the GPU, respectively. Different transformations in the outermost for-loops will result in different features extracted from the innermost assignment statement, leading to varying throughput predictions and learning processes in the cost model. Overall, our method considers various optimizations of the outermost loops during the optimization process. By extracting features from the innermost non-loop statements, we train a learnable cost model on the GPU to automatically optimize the full tensor program. We train the cost model and mainly target the data parallel tensor programs which consist outermost loop-nest and some innermost assignment statement. We use the ATFormer to predict the score of each non-loop innermost statement. An example are shown as follows:\n\n```python\nfor i in range(10):\n  for j in range(10):\n      B[i][j] = A[i] * 2                          # Statement B\nfor i in range(10):\n  \tC[i] = B[i][i] - 3                          # Statement C\n```\n\nIn this example, there are two non-loop innermost statement, B and C. We estimate the cost of these tensor programs. The total cost of it is the sum of cost B and cost C. We perform static analysis to extract the features for every non-loop innermost statement and make prediction for them. These predictions, along with features extracted from the entire program, such as the number of *used cached lines*, *used memory*, *reuse distance*, and *arithmetic intensity curves*  are used to estimate the cost of the entire program. In the following code snippet, we utilize an example of matrix multiplication and ReLU function to elucidate the high-level structure and low-level details as mentioned before.\n```python\nfor i in range(512):\n  for j in range(512):\n    for k in range(1024):\n      C[i, j] += A[i, k] * B[k, j]\nfor i in range(512):\n  for j in range(512):\n    D[i, j] = max(C[i, j], 0)\n```\nOne of the possible high-level structure transformation is shown as below:\n```python\nfor i.0 in range(Tile_I0):\n  for j.0 in range(Tile_J0):\n    for i.1 in range(Tile_I1):\n      for j.1 in range(Tile_J1):\n        for k.0 in range(Tile_K0):\n          for i.2 in range(Tile_I2):\n            for j.2 in range(Tile_J2):\n              for k.1 in range(Tile_I1):\n                for i.3 in range(Tile_I3):\n                  for j.3 in range(Tile_J3):\n                    C[...] += A[...] * B[...]\n        for i.4 in range(Tile_I2 * Tile_I3):\n          for j.4 in range(Tile_J2 * Tile_J3):\n            D[...] += max(C[...], 0.0)\n```\nThe following two code snippets demonstrate the low-level details that can be generated based on the aforementioned high-level structure in the hierarchical search space.\n```python\nparallel i.0@j.0@i.1@j.1 in range(256):\n  for k.0 in range(32):\n    for i.2 in range(16):\n      unroll k.1 in range(16):\n        unroll i.3 in range(4):\n          vectorize j.3 in range(16):\n            C[...] += A[...] * B[...]\n  for i.4 in range(64):\n    vectorize j.4 in range(16):\n      D[...] = max(C[...], 0.0)\n```\n```python\nparallel i.2 in range(16):\n  for j.2 in range(128):\n    for k.1 in range(512):\n      for i.3 in range(32):\n        vectorize j.3 in range(4):\n          C[...] += A[...] * B[...]\nparallel i.4 in range(512):\n  for j.4 in range(512):\n    D[...] = max(C[...], 0.0)\n```\nBy examining the provided code snippets, it becomes evident that throughout the entire process of training and prediction in the cost model, there are considerations for common optimizations such as tiling that affect the outermost for-loops. To provide a more comprehensive understanding of our design and optimization methodology, we will showcase detailed code examples in the appendix. These examples will illustrate the process and incorporate common optimizations to enhance readers' comprehension.\n\n**Q2:** In Figure 3, why are the outputs of self-attention and embedding concatenated again? What is the reason for including an embedding layer in every attention block? Additionally, could you confirm whether the attention block utilizes pre- and post-layer norms?\n\n**A2:** Thanks for the comments. In Figure 3, the term \"embedding\" in the attention block refers to a vector representation of the features in an embedding space. The embedding layer is positioned between the features extracted from assignment statements and the first attention block, as illustrated in the following diagram.\n```\nFeature => Embedding Layer => Attention Block => Attention Block => MLP => Score\n```\nThe function of the embedding layer is to project the discrete features from low dimension space to a continuous embedding space for importance measurement.  This enables the cost model to process the feature and consider the importance of assignment statement in subsequent self-attention mechanisms. Therefore, the embedding layer appears only once in the ATFormer series models. The specific implementation of the embedding layer is described below:\n```python\n      self.embedding_layer = torch.nn.Sequential(\n            torch.nn.Linear(in_dim, hidden_dim),\n            torch.nn.ReLU(),\n            torch.nn.Linear(hidden_dim, hidden_dim),\n            torch.nn.ReLU(),\n        )\n```\nThe embedding layer consists of two consecutive MLP layers followed by a non-linear activation function (ReLU). Therefore, the term \"embedding\" in the Figure 3 does not refer to the embedding layer. We apologize for the confusion caused by the figure and sincerely appreciate the reviewer for pointing out this concern. In the final revised version, we will re-adjust the content of Figure 3 to enhance clarity and facilitate a better understanding of our structure by the readers. In this case, we do not concatenate the output of the embedding layer with the output of self-attention. Instead, we add the output of the embedding layer with a normalization layer to the output of self-attention in the first attention block. In the second attention block, we first pass the output of the first attention block through a normalization layer. Then, we add it to its own self-attention and further pass it through another normalization layer. This addition follows the same principle as the residual addition in standard transformer encoder. Yes. in our ATFormer design, each attention block incorporates two normalization layers: one located before the self-attention mechanism and another positioned after it. The purpose of the normalization layer is to stabilize the input distribution of the forward data and ensure stable gradient propagation during training. \n\n**Q3:** Please provide the latency values for BERT-large and GPT models in Figure 6? It seems that ATFormer offers little improvement compared to MHA in some cases.\n\n**A3:** Figure 6 illustrates the end-to-end network optimization using the traditional learning. In comparison to MHA, ATFormer achieves acceleration factors of 1.09x, 1.11x, 1.06x, 1.21x, 1.04x, and 1.05x for AlexNet, VGG-16, MobileNet-V2, ResNet-18, ResNet-50, and Bert-Tiny, respectively. Regarding Bert-large and GPT-2-large models, ATFormer achieves improvements of **1.11x** and **1.13x**, respectively, compared to MHA.  For further optimization results, please refer to the table provided below: \n| Cost Model| Bert-large (ms) | GPT-2-large (ms) |\n| ----------- | --------------- | ---------------- |\n| XGBoost     | 34.56           | 279.11           |\n| LSTM        | 31.22           | 269.58           |\n| MHA         | 29.98           | 262.74           |\n| ATFormer-1L | 29.49           | 251.42           |\n| ATFormer-2L | 27.01           | 232.51           |\n\n**Q4:** It seems that ATFormer with 2 attention blocks performs worse than ATFormer-1L on the dataset with a size of 300 points in Figure-7. Can you provide insights into the reasons behind this performance difference?\n\n**A4:** Thanks for the comments. Figure 7 presents the transfer learning capabilities of different cost models on the TenSet-50/100/200/ 300/500 datasets. It is discernible that the augmentation in inference performance differs among the various cost models when confronted with a relatively constrained offline dataset (TenSet-50).  Subsequently, as the offline dataset's sample count increases, the transfer learning capabilities demonstrated by the distinct cost models progressively evolve.Particularly, when examining TenSet-300, it becomes evident that the transfer learning prowess of ATFormer-2L lags behind that of ATFormer-1L. The underlying cause lies in the fact that ATFormer-1L possesses a smaller parameter size in comparison to ATFormer-2L, resulting in the latter model experiencing underfitting after training on a constrained offline dataset. Nonetheless, as the number of training samples within the dataset grows, the learning capability of ATFormer-2L demonstrates continuous improvement, eventually surpassing ATFormer-1L in terms of optimization performance. Overall, ATFormer-M showcases the most stable performance among all the cost models during the transfer learning. It provides strong evidence for the superiority of the mask mechanism when training attention-based cost models during the process of operator tuning.\n\n**Q5:** In Table-2, it appears that ATFormer reduces search time by utilizing transfer learning, but this comes at the cost of increased kernel execution latency. Does this ultimately lead to a worse overall result?\n\n**A5:** Thanks for the comments. In Table-2, the term \"transfer learning\" refers to the training of cost models such as XGBoost, LSTM, MHA, ATFormer-1L, and ATFormer (2L) on the Tenset-500 offline dataset, followed by subsequent online updates of the pre-training. On the contrary, \"traditional learning\" indicates the direct training of these cost models from scratch using online datasets. It is worth noting that the number of trials required during the tuning (search) process differs between these two approaches. The table clearly demonstrates a significant reduction in the tuning time taken for the networks to converge, for example, from 3028 to 654 seconds as for BERT-base model. In the \"Speedup\" metric, it represents the acceleration achieved at the latency level when tuning a network end-to-end with the same trials, such as a 1.39x speedup.  The \"time\" metric indicates the substantial acceleration in achieving approximate inference latency with both traditional learning and transfer learning, for instance, a 4.97x acceleration. More details about these two metrics can be found in line 432-435. Therefore, using transfer learning only accelerates the convergence time during network tuning and does not increase the cost of kernel execution or lead to worse results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1)Extensive experiments. The results can be useful. \n2) An interesting approach. The proposed approach uses two parts of feature information to capture high-level and low-level representations separately: coarse-grained operator embedding and fine-grained statement.", "reasons_to_reject": "The code generation templates are overly simple. The final prediction/tuning objective only considers the innermost non-loop statements. There is no optimization like loop tiling in the code generation task, which is shown to be important for code optimization of DNN workloads.", "questions_to_authors": "1.\tIn Figure 3, why are the outputs of self-attention and embedding concatenated again? What is the reason for including an embedding layer in every attention block? Additionally, could you confirm whether the attention block utilizes pre- and post-layer norms?\n2.\tPlease provide the latency values for BERT-large and GPT models in Figure 6? It seems that ATF offers little improvement compared to MHA in some cases.\n3.\tIt seems that ATFormer with 2 attention blocks performs worse than ATFormer-1L on the dataset with a size of 300 points in Figure-7. Can you provide insights into the reasons behind this performance difference?\n4.\tIn Table-2, it appears that ATF reduces search time by utilizing transfer learning, but this comes at the cost of increased kernel execution latency. Does this ultimately lead to a worse overall result?", "paper_topic_and_main_contributions": "This paper provides a comprehensive analysis of auto-tuning challenges in code generation. Instead of relying on a statistical model like XGBoost, the proposed techniques leverage an attention-block-based model to capture the implicit long-context relationships and representations in tensor programs. Compared to other DNN-based algorithms, such as LSTM, the attention mechanism executes in parallel and requires less training time.\n\nMoreover, the paper showcases the applicability of cross-device transfer learning, effectively transferring learned features from offline static datasets on one hardware platform to another new platform. This approach enables users to expedite the optimization/searching stage by reusing a cost model.\n\nThe study includes a detailed comparison experiment involving existing algorithms (XGBoost, LSTM, and MHA) on two NVIDIA GPU architectures: Turing and Ampere. These experiments encompass widely used classical CV and NLP models like VGG, ResNet, BERT, and GPT. Additionally, the paper conducts a transfer learning evaluation to gauge the performance of the transferred features during tuning for tensor program generation tasks.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces ATFormer, a cost modeling and tensor program optimization framework for enhancing deep neural network efficiency. By framing the tuning problem as an optimization task in a hierarchical search space, ATFormer employs attention mechanisms to capture global and long-range dependencies in features. This model is assessed both traditionally and through transfer learning. 

Most reviewers asked for additional explanation/results related to code generation of the approach. Initially, reviewers interpreted the work to optimize only innermost loops, limiting its applicability. However, the authors provided extensive feedback and better explanation for the method (and example output) in the rebuttal and convinced the reviewers that the approach is a more general framework. I suggest that the authors try to refactor the paper to make this more clear and provide some more interesting examples upfront (taken from the rebuttal discussion) to highlight to what degree this work can perform optimization for tensor programs.
Decision: Accept-Main