CoDSA: A Hybrid Tensor Compilation Framework Integrating Coordinate Descent with Dynamic Simulated Annealing

Abstract

The deployment of Deep Learning models on heterogeneous hardware necessitates optimized tensor programs to utilize computational resources. While auto-tuning frameworks like Ansor employ evolutionary search strategies, they suffer from limitations, including slow convergence, memory consumption, and a tendency to stagnate in local optima within high-dimensional, non-convex search spaces. To address these challenges, this paper proposes CoDSA (Coordinate Descent with Dynamic Simulated Annealing), a hybrid auto-tuning framework that integrates an enhanced Droplet Search algorithm into the Ansor ecosystem. We introduce three improvements to the droplet coordinate descent strategy: 1) a Dynamic Simulated Annealing mechanism to escape local optima; 2) an Adaptive Step-Size Control modulated by an exploration factor to balance global traversal and local refinement; and 3) a Diversity Control mechanism utilizing cosine similarity to enforce orthogonal exploration. Experiments on NVIDIA GPUs using Vision Transformer, DeiT, and Swin Transformer models demonstrate that CoDSA outperforms the Ansor framework. Specifically, for the ViTBase model, our method reduces search time by 37.2% and memory usage by 53.1%, while improving inference latency by up to 26.5% on complex architectures like Swin-Small. This work bridges the gap between the global exploration capability of evolutionary methods and the rapid exploitation efficiency of coordinate descent.