arXiv:2604.12891v1 [cs.LG] 14 Apr 2026
TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning CHAOYAO SHEN and LINFENG JIANG, Southeast University, China YIXIAN SHEN, University of Amsterdam, Netherlands TAO XU, Southeast University, China GUOQING LI∗ , Southeast University & Shandong Yunhai Guochuang Innovative Technology Co., Ltd., China ANUJ PATHANIA, University of Amsterdam, Netherlands ANDY D. PIMENTEL, University of Amsterdam, Netherlands MENG ZHANG∗ , Southeast University, 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.8× and 12.48× faster tuning time, and 1.20× and 1.13× lower inference latency, respectively, compared to Tenset-MLP. CCS Concepts: • Computing methodologies → Machine learning; • Software and its engineering → Compilers. Additional Key Words and Phrases: Deep learning compiler, Tensor program, Continual learning, Active learning, Cost model
1
Introduction
Deep neural networks (DNNs) have achieved remarkable success in domains such as computer vision [17, 25, 40, 41, 49, 63], natural language processing [39, 40, 45, 57], and autonomous driving [8, 47], driving rapid innovation in AI hardware accelerators [2]. Specialized platforms such as NVIDIA GPUs [14], Google TPUs [29], and GraphCore IPUs [28] meet the growing computational demands of these models. Deep learning (DL) compilers [22] have emerged as a critical bridge between high-level DNN models and low-level hardware execution. Traditionally, tensor operators are implemented using manually optimized kernel libraries, such as NVIDIA cuDNN [11] and Intel OneDNN [26], which are finely tuned for specific hardware. While these libraries offer excellent performance, they are often tightly ∗ Corresponding author.
New Paper, Not an extension of a Conference Paper. Authors’ Contact Information: Chaoyao Shen, [email protected]; Linfeng Jiang, [email protected], Southeast University, Nanjing, Jiangsu, China; Yixian Shen, University of Amsterdam, Amsterdam, Netherlands, [email protected]; Tao Xu, Southeast University, Nanjing, Jiangsu, China, [email protected]; Guoqing Li, Southeast University & Shandong Yunhai Guochuang Innovative Technology Co., Ltd., Nanjing, Jinan, Jiangsu, Shandong, China, [email protected]; Anuj Pathania, University of Amsterdam, Amsterdam, Netherlands, [email protected]; Andy D. Pimentel, University of Amsterdam, Amsterdam, Netherlands, [email protected]; Meng Zhang, Southeast University, Nanjing, Jiangsu, China, [email protected]. Manuscript submitted to ACM
1
2
C. Shen et al.
coupled to particular architectures, making them difficult to adapt to new or diverse platforms. Moreover, they struggle to support emerging or non-standard operators, limiting their applicability in rapidly evolving model architectures. To address the limitations, search-based DL compilers [4, 9] have been developed to automatically generate and optimize tensor programs across a variety of hardware backends. Unlike static kernel libraries, these compilers explore the space of schedule transformations and select implementations tailored to the target device. A key component of this process is the cost model, which predicts the performance (e.g., latency or throughput) of candidate tensor programs and guides the search process effectively. Current search-based DL compilers often train cost models online [3, 62], by iteratively sampling candidate tensor programs, evaluating their hardware performance, and updating the model during auto-tuning. While this dynamic feedback loop enables hardware-specific adaptation, it comes at a high cost: each program must be compiled and executed during search, which significantly increases tuning time and demands extensive computing resources. The release of large-scale tensor program datasets, such as Tenset [61], has made it feasible to shift cost model training offline. These datasets provide millions of tensor programs with performance measurements across hardware backends, enabling the development of predictive models without requiring real-time hardware profiling. Motivated by this, recent work has increasingly focused on offline training of cost models [46]. In this paradigm, the cost model is trained once on pre-collected data from a specific hardware platform and can then be used to predict the performance of unseen tensor programs during auto-tuning. To further reduce hardware-specific data collection, researchers have explored transfer learning strategies [16, 38], which adapt cost models pre-trained on one or more source platforms to new target devices by leveraging shared compiler optimization patterns and applying platform-specific fine-tuning. This combination of offline training and transfer learning reduces tuning time and resource usage while improving compiler adaptability across diverse hardware environments. Although DL compilers equipped with offline-trained cost models demonstrate strong performance in terms of tuning quality and efficiency, several practical and architectural challenges remain. First, collecting a sufficiently large set of tensor programs for each hardware platform is highly time-consuming. In our measurements, data collection on an Intel i7-12700F CPU requires approximately 40 days, while the same process on an NVIDIA GeForce RTX 3080Ti GPU takes over 60 days, posing a major bottleneck for scalability. Second, the architecture of the cost model fundamentally limits the accuracy and efficiency of performance prediction. Early approaches relied on XGBoost [9] and MLPs [61], which have since evolved to more expressive models such as LSTMs [5] and Transformers [37]. However, LSTMs suffer from limited parallelism and slow convergence, while Transformers, despite their ability to capture long-range dependencies, impose high computational and memory overhead when applied to long sequences of schedule tokens. These architectural limitations prompted us to design a cost model that can efficiently capture sequence dependencies while maintaining a lightweight approach. Third, while transfer learning has emerged as a popular technique for enabling cross-platform generalization [59, 61], most existing methods are restricted to one-to-one adaptation: knowledge is transferred from a single source hardware to a single target hardware. This restricts their ability to accumulate generalized scheduling knowledge across diverse platforms. Recent work such as MTL-TLP [55] attempts to overcome this limitation using multi-task learning to jointly learn from multiple source platforms. While promising, this approach introduces significant scalability issues: the number of trainable parameters and associated computational cost grow proportionally with the number of hardware targets, and the framework requires simultaneous access to data from all source platforms during training, limiting practicality. These limitations have prompted us to develop more flexible multi-platform learning frameworks to efficiently utilize knowledge from multiple source platforms without introducing excessive computational costs. Manuscript submitted to ACM
TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning
3
To address the challenges of existing DL compilers across hardware platforms, namely the high cost of data collection, inefficiencies in cost model design, and suboptimal generalization across hardware, we propose TCL, a fast, data-efficient, and transferable tensor compiler framework. At the data level, it incorporates a principled active learning strategy, the RDU Sampler, which selects the most informative tensor programs by jointly considering representativeness, diversity, and uncertainty. This dramatically reduces the number of performance measurements required, achieving competitive or even superior tuning quality with as little as 10% of the data. To enhance prediction efficiency, TCL employs a new cost model architecture that is both lightweight and expressive, built upon the Mamba Block [19], a state-space sequence model capable of capturing long-range dependencies with linear-time complexity. In contrast to Transformer-based models, which incur O (𝑛 2 ) time and memory complexity due to their attention mechanisms, Mamba reduces this overhead to O (𝑛) by leveraging structured state-space representations. This design yields comparable prediction accuracy while significantly improving training and inference efficiency, making it well-suited for compiler deployment across diverse hardware platforms. For cross-hardware adaptation, TCL introduces a continual knowledge distillation framework that enables progressive knowledge accumulation from multiple platforms. By decoupling shared knowledge representation (the hardware knowledge base) from hardware-specific adaptation (the hardware active column), the framework avoids the parameter growth and data simultaneity constraints inherent in multi-task learning, enabling more scalable and modular transfer across platforms. Our key contributions are as follows: • We introduce the RDU Sampler, an active learning algorithm that jointly optimizes representativeness, diversity, and uncertainty to select high-value tensor programs. Using only 10% of the dataset, it attains performance close to or better than models trained on the full dataset, significantly reducing data collection overheads. • We design a new cost model architecture based on the Mamba Block, which efficiently captures long-range dependencies in schedule sequences while significantly reducing computational and memory complexity, from O (𝑛 2 ) in Transformer-based models to O (𝑛). This lightweight architecture enables faster training and improved prediction accuracy with fewer parameters. • We propose a continual knowledge distillation framework for scalable cross-hardware adaptation. It consists of a shared hardware knowledge base that accumulates information from multiple source platforms, and a hardware active column that transfers this knowledge to new hardware targets. This progressive design avoids the parameter explosion and simultaneous data dependency issues of conventional multi-task learning. • We open-source a large-scale dataset of tensor programs, collected over several months on an Intel i7-12700F CPU and an NVIDIA GeForce RTX 3080Ti GPU1 . Extensive experiments on this dataset demonstrate that TCL substantially reduces tuning time and improves inference latency compared to state-of-the-art baselines. 2
Related Works
Active learning [42, 53] reduces the annotation workload by selectively labeling the most informative samples, enabling models trained on these subsets to achieve high performance with less data. While active learning has been extensively studied in classification settings [13, 23, 31, 43], its application to regression tasks remains relatively limited [6, 15]. In the context of DL compilers, offline cost model training typically requires large-scale tensor program datasets. Although effective, recollecting such datasets for each new hardware platform incurs significant time and computational costs. To address this challenge, recent studies have investigated active learning and sampling strategies within DL compilers. ALT [54] was the first to integrate active learning into a DL compiler, effectively reducing tuning time in online cost 1 https://github.com/booker0415/Large-Scale-Tensor-Program-Dataset-on-RTX-3080-Ti-and-Intel-i7-12
Manuscript submitted to ACM
4
C. Shen et al.
model training. BALTO [7] proposed a biased-diversity selection method to accelerate offline training. CDPMM [24] introduced a clustering-based sampling approach that selects representative tensor programs, improving adaptation to new hardware and enhancing generalization. Cost model architecture in DL compilers must balance prediction accuracy with computational efficiency. Early efforts relied on analytical heuristics; e.g., Halide [33] estimated performance using data reuse and FLOP counts for loop transformations. Subsequent approaches adopted machine learning models: AutoTVM [10] and Ansor [60] employed XGBoost to predict relative performance rankings, while Tenset [61] showed that a combination of MLP and RankLoss yielded superior accuracy. More recent works explored deep learning-based architectures. TIRAMISU [4] applied LSTMs to model sequential schedule behavior, whereas One-Shot [37], TLP [55], and CDPMM [24] adopted Transformer-based encoder-decoder structures. Despite their modeling power, LSTMs suffer from limited parallelism, and Transformers incur quadratic complexity with respect to sequence length, making them costly for large-scale tuning scenarios. Cross-hardware generalization is increasingly critical in DL compiler optimization, as tuning time is dominated by hardware measurements, which often become the primary bottleneck [35, 44, 51]. Offline-trained cost models alleviate this by removing the need for on-device profiling, but they are tied to specific platforms and require costly data recollection and retraining for new hardware. To address this, transfer learning techniques [32, 52, 56] have been explored to reuse knowledge from source platforms and accelerate model convergence on target devices. Moses [59], grounded in the lottery ticket hypothesis, uses domain adaptation to disentangle hardware-aware and hardwareagnostic features for better cross-device generalization. Verma et al. [48] propose an attention mechanism that embeds transferable hardware-specific signals into cost models. TLP [55] employs a multi-task learning strategy with shared parameters to learn general optimization patterns and task-specific branches to capture hardware-specific nuances. While these methods improve cross-platform tuning efficiency, challenges remain in balancing generalization and scalability. To bridge these gaps, we propose TCL, a unified framework that integrates: (1) the RDU Sampler for data-efficient tensor selection, (2) a Mamba-based cost model for fast and accurate prediction, and (3) a continual knowledge distillation mechanism that supports scalable and effective cross-platform generalization without multi-tasking overhead.
3
Preliminaries and Overview
3.1
Preliminaries
In this section, we define some key terms and concepts used throughout the paper. These definitions establish a foundation for the proposed methods and analyses. Network and Subgraph. A DNN is defined with a specific input size, such as batch size and image input dimensions (e.g., ResNet-18, MobileNetV1). A DL compiler partitions the network into subgraphs based on specific rules, where each subgraph represents a part of the original network. Hardware Platform. The devices on which DNNs are deployed and executed are referred to as hardware platforms. Here, devices with different hardware architectures are considered distinct hardware platforms. We do not transfer knowledge between CPU and GPU in this paper. Assignment. The execution of a subgraph on a hardware platform is termed an assignment. A DNN can be divided into multiple subgraphs, and when these subgraphs are executed on a hardware platform, they result in multiple assignments. For example, ResNet-18 partitions into 18 subgraphs in Ansor, leading to the creation of 18 assignments. Manuscript submitted to ACM
TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning
5
Tensor programs 1
+
Subgraph
Hardware Platform
Schedule primitives 1 =
Input Size 1, 3, 224, 224
Assignment
Conv2d 3 , 64, 7, 7
BN 1, 1, 1, 64
Schedule primitives 2 Tensor programs 2
Relu
Output Size 1, 64, 112, 112
Fig. 1. Examples of two tensor programs for Conv + BN + ReLU subgraph in ResNet-18 on Nvidia Tesla T4
Schedule Primitive. Schedule primitives are transformations used in DL compilers to optimize the execution of computational graphs across platforms, enhancing performance metrics such as latency, memory usage, and computational efficiency. Commonly used schedule primitives include SP (split), RE (reorder), and FU (fuse), among others. Tensor Program. A tensor program is a low-level implementation of tensor operations that specifies how computations are performed on a hardware platform. We present an assignment for a subgraph (Conv + BN + ReLU) from ResNet-18 on the Nvidia Tesla T4 hardware platform in Fig. 1, along with two variants of the schedule primitives and tensor programs for this assignment. 3.2
Overview
TCL aims to leverage cost models trained offline on multiple source hardware platforms to facilitate the offline training of cost models on target hardware platforms. Fig. 2 illustrates the flowchart of TCL. TCL accepts DNN models (e.g., PyTorch[34], TensorFlow[1], Keras[12], and Caffe[27]) developed in various DL frameworks. It then performs graphlevel optimizations (e.g., operator fusion, constant folding, and data layout transformation) to output a computational subgraph. The granularity of the subgraph is the same as that of the Ansor, and the DNN is converted into a small subgraph using the operator fusion algorithm of Relay[36]. This subgraph is processed by a search algorithm to generate schedule primitive sequences, which are used for training the cost model. Subsequently, TCL operates in two phases: training and inference. During the training phase, the offline training of the cost model on the target hardware platform relies on a continuous knowledge distillation framework that accumulates knowledge from cost models trained offline on multiple Manuscript submitted to ACM
6
C. Shen et al. Deep Learning Frameworks Inference DNN Model
Train Graph-level Optimization Seach Algorithm Schedule primitive sequence
Computational Subgraph
TCL
Top-k schedule primitive sequence
Schedule primitive sequence
RDU Sampler
Mamba-based cost model on others source hardware (offline training)
Knowledge distillation
Selected schedule primitive sequence
Mamba-based cost model on target hardware (offline training)
Code Generation
Extract accumulated knowledge
Continuous Knowledge Distillation Framework
Tensor Programs
Target Hardware
Fig. 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.
source hardware platforms. The trained cost model on each source device accumulates knowledge into the continuous knowledge distillation framework through knowledge distillation, akin to how the human brain accumulates old knowledge to facilitate the learning of new knowledge. When training the cost model for the target hardware platform, in addition to using the continuous knowledge distillation framework, it is also necessary to use the RDU Sampler for data sampling to save data collection time. In the inference phase, schedule primitive sequences generated by the search module predicts the performance of the tensor program using the cost model. The cost model selects the best K optimal schedule primitive sequences, generates the tensor program, and deploys it on the target hardware platform. The cost model effectively identifies the relationships between schedule primitive sequences through an encoder-decoder designed based on the Mamba block, focusing on sequences that significantly impact performance. 4
RDU Sampler
The RDU Sampler is a key component of TCL, designed to select the most informative and valuable samples from the hardware platform. This process reduces data collection time during offline training while ensuring the predictive accuracy of the cost model. Since the performance of the cost model relies heavily on the quality of the samples chosen by the RDU Sampler, it is essential to carefully design the sampling criteria. In our RDU Sampler, we develop a high-performance sampling algorithm based on three criteria: representativeness, diversity, and uncertainty. Representativeness should accurately reflect the characteristics of the entire dataset distribution, ensuring that the sampled data retains a similar class or feature distribution to the original dataset. Fig. 3 shows the distribution of Manuscript submitted to ACM