Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning

TL;DR

This work tackles the inefficiency of search-based tensor program tuning caused by slow learned cost models and cross-platform online unawareness. It introduces Pruner, a Draft-then-Verify exploration mechanism with a Latent Schedule Explorer for rapid draft generation and a Pattern-aware Cost Model for accurate verification, complemented by MoA-Pruner, a momentum online adaptation strategy for cross-platform transfer. Across three GPU platforms, the approach achieves substantial speedups over state-of-the-art baselines in both online and offline tuning (e.g., average online speedups of $2.6\times$ for Pruner and $4.82\times$ for MoA-Pruner vs Ansor, and offline gains around $4.75\times$–$4.05\times$ vs TenSet/TLP with an additional $4.08\times$ over MetaSchedule on TensorCore). Implemented in the TVM framework, Pruner demonstrates strong end-to-end performance gains, robust single-operator results, and favorable compilation costs, highlighting its practical impact for efficient tensor-program tuning.

Abstract

Tensor program tuning is essential for the efficient deployment of deep neural networks. Search-based approaches have demonstrated scalability and effectiveness in automatically finding high-performance programs for specific hardware. However, the search process is often inefficient, taking hours or even days to discover optimal programs due to the exploration mechanisms guided by an accurate but slow-learned cost model. Meanwhile, the learned cost model trained on one platform cannot seamlessly adapt online to another, which we call cross-platform online unawareness. In this work, we propose Pruner and MoA-Pruner. Pruner is a "Draft-then-Verify" exploration mechanism that accelerates the schedule search process. Instead of applying the complex learned cost model to all explored candidates, Pruner drafts small-scale potential candidates by introducing a naive Symbol-based Analyzer (draft model), then identifies the best candidates by the learned cost model. MoA-Pruner introduces a Momentum online Adaptation strategy to address the cross-platform online unawareness. We incorporate Pruner into the TVM and conduct extensive experiments on three GPU-based platforms. Results show considerable speedup in schedule search time. In online tuning scenarios, Pruner and MoA-Pruner achieve an average speedup of $2.6 \times$ and $4.82 \times$ compared to Ansor. In offline tuning scenarios, Pruner achieves an average speedup of $4.75 \times$ and $4.05\times$ compared to TenSet and TLP, respectively. Furthermore, Pruner achieves an average speedup of $4.08 \times$ compared to MetaSchedule on TensorCore.

Figures (16)

• Figure 1: The workflow of search-based DLCs. The red dashed box is the optimization workspace of the Pruner.
• Figure 2: The system overview of Pruner, which contains latent schedule explorer and pattern-aware cost model. Momentum online adaptation is only activated for MoA-Pruner in online cost model tuning scenarios.
• Figure 3: An illustrative example of the hardware-aware symbols extraction process for a GEMM-ReLU graph. Some schedule primitives in the schedules generation rules and hardware symbols for some statements are omitted for brevity. Prod or Sum means that the variable equals the product or sum of an array of variables.
• Figure 4: (top) The pipeline of Pattern-aware Cost Model; (bottom) Extraction of temporal dataflow feature.
• Figure 5: Overview of MoA, where $\mathcal{\phi}_{s}$ and ${\Delta}_{t}$ refer to the parameters of the Siamese and gradient of target model, the momentum $m=0.99$. The red arrow means offline pre-train.