TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning
Authors: Chaoyao Shen, Linfeng Jiang, Yixian Shen, Tao Xu, Guoqing Li, Anuj Pathania, Andy D. Pimentel, Meng Zhang
Organizations: Southeast University, China · University of Amsterdam, Netherlands · Southeast University & Shandong Yunhai Guochuang Innovative Technology Co., Ltd., China
Deep learning (DL) compilers rely on cost models and auto-tuning to optimize tensor programs for target hardware. However, existing approaches depend on large offline datasets, incurring high collection costs and offering suboptimal transferability across platforms. In this paper, we introduce TCL, a novel efficient and transferable compiler framework for fast tensor program optimization across diverse hardware platforms to address these challenges. Specifically, TCL is built on three core enablers: (1) the RDU Sampler, a data-efficient active learning strategy that selects only 10% of tensor programs by jointly optimizing Representativeness, Diversity, and Uncertainty, substantially reducing data collection costs while maintaining near-original model accuracy; (2) a new Mamba-based cost model that efficiently captures long-range schedule dependencies while achieving a favorable trade-off between prediction accuracy and computational cost through reduced parameterization and lightweight sequence modeling; and (3) a continuous knowledge distillation framework that effectively and progressively transfers knowledge across multiple hardware platforms while avoiding the parameter explosion and data dependency issues typically caused by traditional multi-task learning. Extensive experiments validate the effectiveness of each individual enabler and the holistic TCL framework. When optimizing a range of mainstream DL models on both CPU and GPU platforms, TCL achieves, on average, 16.8x and 12.48x faster tuning time, and 1.20x and 1.13x lower inference latency, respectively, compared to Tenset-MLP.
Figures & tables
Figure 1. Examples of two tensor programs for Conv + BN + ReLU subgraph in ResNet-18 on Nvidia Tesla T4
Figure 2. Overview of the TCL cross-hardware DL Complier. The training phase offline-trains the cost model for the target hardware, leveraging models from other hardware. The inference phase integrates the trained model into the DL compiler’s end-to-end pipeline. The light black blocks represent the modules based on Ansor, the light pink blocks represent our proposed TCL modules, and the light yellow blocks represent the specific components within TCL.
Figure 3. Distribution of operator types for assignments on Intel Xeon E5-2673 in the Tenset dataset. This dataset covers 2,308 assignments from multiple deep learning models.
Figure 4. Normalized latency score frequency distribution for tensor programs on Intel Xeon E5-2673 in the Tenset dataset. The normalized latency score refers to the ratio between the latency of tensor programs generated using different schedule primitives on the hardware platform and the minimum latency achieved for that task across all schedule primitives.
CPU Features
GPU Features
Cores
Cuda Cores
Threads
Base Clock
Turbo frequency
Boost Clock
Boost Frequency
Memory clock
L1 Cache
Memory size
L2 Cache
Memory bus width
Table 1. CPU and GPU Feature Parameters
Figure 5. Model architecture description of Mamba-based cost model.
Figure 6. Overview of the proposed Continuous Knowledge Distillation (CKD) framework for cross-hardware cost model training. For each new hardware platform (e.g., CPU, GPU), an active column (AC) learns platform-specific knowledge from new tensor program traces while the knowledge base (KB) remains frozen to preserve prior knowledge. After the continual learning (CL) phase, knowledge from the AC is selectively distilled into the KB via a gated distillation mechanism and constrained by Elastic Weight Consolidation (EWC) to prevent catastrophic forgetting. This process enables progressive, scalable adaptation across diverse hardware targets without requiring simultaneous multi-platform data or expanding model size. Both the AC and KB share the same Mamba-based architecture (encoder, Mamba block, decoder).
T4
E5-2673
nlayer
dstate
expand
dconv
Top-1 Score
Top-5 Score
Top-1 Score
Top-5 Score
1
32
1
4
0.856
0.91
0.8913
0.9605
1
16
1
4
0.8848
0.954
0.8530
0.9502
1
8
1
4
0.8891
0.9532
0.9050
0.9639
2
8
1
4
0.8512
0.9314
N/A
N/A
3
8
1
4
0.8458
0.9061
N/A
N/A
Table 2. Experiments on different hyperparameter combinations of Mamba-based cost model
Hardware platforms
Tenset
TLP
Ours
Top-1 Score
Top-5 Score
Size
Top-1 Score
Top-5 Score
Size
Top-1 Score
Top-5 Score
Size
Intel Platinum 8272CL
0.8748
0.9527
0 .92MB
0.9194
0.9710
1 .32MB
0.9132
0.9592
0.35MB
Intel E5-2673 v4
0.8332
0.8977
0.8941
0.9633
0.9050
0.9639
AMD EPYC 7452
0.8510
0.9175
0.9055
0.9494
0.9058
0.9619
ARM Graviton2
0.7799
0.9049
0.8207
0.9226
0.8300
0.9264
Intel i7-12700F
0.8807
0.9475
0.9068
0.9537
0.9229
0.9691
Table 3. Comparison of Cost Models on Different Hardware Platforms
Figure 7. Comparison Across Different Sampling Rates on CPU and GPU Platforms.
T4
E5-2673
Top-1 Score
Top-5 Score
Top-1 Score
Top-5 Score
ALT
0.8668
0.9305
0.8521
0.9359
BALTO
0.8466
0.8933
0.8619
0.9199
Ours
0.8761
0.9318
0.8806
0.9512
Table 4. Comparison of Sampling Methods on CPU and GPU Platforms
Tensor program optimization is essential for modern machine learning systems, but its search space is enormous. Existing auto-schedulers reduce measurement cost with learned cost models, yet they usually evaluate each candidate as a static code snapshot, ignoring the schedule trajectory that produced it. This makes them insensitive to action dependencies and vulnerable to superficial code variations. We propose a world-model-inspired evaluator that models schedule evaluation as action-conditioned latent dynamics over program states. Starting from the initial program, it rolls out scheduling actions in a continuous latent space with a lightweight transition model, avoiding expensive AST mutation and repeated code encoding. The final dynamic representation is combined with action and hardware features to rank candidates. Implemented in TVM AutoScheduler, our method improves representative-subgraph latency over Ansor by 1.37× on GPU and 1.54× on CPU under the same 64-trial budget. It also matches Ansor-10K within 2.2% geometric mean using 10× fewer measurements, and accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61×/3.67× geometric mean.
Haolin Pan, Lianghong Huang, Xvlin Zhou +2
Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences, Hangzhou, China · Institute of Software, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
With the widening gap between compute and memory operation latencies, data movement optimizations have become increasingly important for DNN compilation. Current optimizations such as layout transformations and operator fusion only target a subset of tensor operators and consequently miss important opportunities for reducing data movement in contemporary DNN workloads, including large language models. We introduce VTC, a novel tensor compilation framework that for the first time eliminates all unnecessary data movement by targeting the full spectrum of data movement operators. VTC proposes the concept of virtual tensors to track data movement between compute operators via index mappings rather than expensive physical data transfers to and from global memory, which can seamlessly interoperate with existing computation kernels and handle arbitrary tensor operator compositions. We also introduce a novel data movement elimination algorithm to automatically identify a profitable virtual tensor creation strategy. Evaluation on a variety of DNNs shows that VTC can outperform existing ML compilers by up to 1.93x (1.28x on average) on NVIDIA GPUs with up to 60% (17.5% on average) inference memory savings.
Muyan Hu, Ahan Gupta, Jiachen Yuan +7
University of Illinois Urbana-Champaign · Georgia Institute of Technology · Microsoft
We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.