Search papers, labs, and topics across Lattice.
This paper introduces TCL, a compiler framework for optimizing tensor programs across diverse hardware using continual learning. TCL employs a Representativeness, Diversity, and Uncertainty (RDU) sampler for data-efficient active learning, a Mamba-based cost model for capturing schedule dependencies, and a continuous knowledge distillation framework for cross-hardware transfer. Experiments show TCL achieves significant speedups in tuning time (16.8x on CPU, 12.48x on GPU) and reduced inference latency (1.20x on CPU, 1.13x on GPU) compared to Tensor-MLP.
Forget massive offline datasets: TCL slashes tensor program optimization time by 16x while *improving* inference latency, thanks to a Mamba-powered cost model and continual learning.
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.
