TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning

TL;DR

TCL is a new compiler framework that uses continual learning and active data selection to optimize tensor programs efficiently across various hardware platforms, reducing costs and improving transferability.

Contribution

It introduces three core components: an active learning sampler, a lightweight cost model, and a knowledge distillation framework for cross-platform transfer learning.

Findings

Achieves 16.8x faster tuning on CPU and 12.48x on GPU.

Reduces data collection costs by selecting only 10% of programs.

Improves inference latency by 1.20x on CPU and 1.13x on GPU.

Abstract

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.