Using Task Graph Caching to Accelerate TVM Code Generation

Deep Learning (DL) models are at the core of a growing number of applications, making fast, low-latency execution across diverse device architectures both a critical requirement and a challenge. DL compilers, such as TVM, address this challenge by automatically translating high-level models into optimized low-level code that effectively exploits device architectures. However, search-space-based algorithms face difficulties in exploring the vast optimization sequence space, often resulting in lengthy compilation times that significantly impact the design cycle. This paper introduces the Task Graph Caching (TGC) algorithm, which aims to reduce the high compilation time while preserving the quality of the code generated by TVM. In particular, TGC enhances TVM auto-tuning by exploiting the fact that similar DL subgraphs appear both within and across models, thus enabling optimization sequences discovered in past compilations to guide future executions. To achieve this, TGC introduces a cache structure that stores high-performance optimization sequences found in previous TVM executions. This information is then used to seed the population of TVM evolutionary search, avoiding redundant exploration of the optimization space and accelerating convergence. Experimental results on twelve DL models show that TGC can significantly speed up the search for efficient optimization sequences, reducing auto-tuning time by up to 2.89 × for Ansor and 3.13 × for MetaSchedule on CPU. Moreover, TGC reduces auto-tuning time by up to 3.25 × for MetaSchedule on GPU. Furthermore, on average, TGC can maintain the inference time achieved by the default TVM, making it a promising solution for accelerating the compilation of DL models.