The search space for loop optimizations is large, with a single matrix multiplication operator potentially having billions of valid schedule configurations when accounting for tile sizes, vectorization factors, and loop reordering permutations. If benchmarking a single configuration on physical hardware takes even a fraction of a second, evaluating the entire space is computationally impossible.
To navigate this massive set of possibilities efficiently, we cannot rely on hardware measurements for every candidate. Instead, we use a cost model. A cost model serves as a mathematical proxy for the hardware. It predicts the performance of a specific loop schedule without actually running the code. This allows the auto-tuner to simulate millions of configurations per second, filtering out poor candidates and identifying promising ones for actual verification.
The fundamental goal of the cost model is to approximate the execution time of a generated kernel. If we define a schedule configuration as and the actual hardware execution time as , the cost model attempts to minimize the error:
In a typical auto-tuning pipeline, the search algorithm proposes a batch of candidate schedules. The cost model evaluates these candidates and assigns a score or predicted latency to each. Only the candidates with the best predicted scores are sent to the hardware for benchmarking. This dramatically reduces the tuning time by filtering out inefficient schedules early in the process.
There are two primary categories of cost models used in machine learning compilers: analytical models and learned models.
Early compilers relied heavily on analytical modeling. These models use predefined formulas based on hardware specifications. An analytical model works by counting operations and estimating latency based on instruction costs.
For example, a simple analytical model for a CPU might calculate cost as:
Where represents the cost in clock cycles for floating-point operations, memory loads, and cache misses respectively.
While analytical models are extremely fast to evaluate, they often fail to capture the complexity of modern accelerators. Modern GPUs and TPUs have complex behaviors such as out-of-order execution, bank conflicts, and instruction pipelining that are difficult to express in a static formula. Consequently, a schedule that looks efficient analytically might perform poorly in practice due to hardware-specific bottlenecks like shared memory bank conflicts.
Modern frameworks like TVM and Halide have moved toward learned cost models. Instead of relying on a fixed formula, these systems use machine learning to predict the performance of other machine learning workloads. This approach treats performance prediction as a regression or ranking problem.
A learned cost model typically uses a lightweight statistical method, such as XGBoost (Gradient Boosted Trees) or a small Multi-Layer Perceptron (MLP). The workflow operates as follows:
This approach is effective because it adapts to the specific hardware. If you tune a model on an NVIDIA A100, the cost model learns the performance characteristics of the A100. If you switch to a mobile ARM CPU, the cost model retrains and learns the distinct characteristics of that architecture.
To train a model, the Abstract Syntax Tree (AST) of the loop nest must be converted into a feature vector. Common features extracted by ML compilers include:
The diagram below illustrates the flow of data from the search space through the cost model and into the hardware measurement phase.
The iterative workflow of a learned cost model showing how hardware measurements refine predictions over time.
A critical challenge in auto-tuning is the balance between exploration and exploitation.
If the tuner only exploits, it might get stuck in a local minimum because the cost model is inaccurate in unexplored regions. If it only explores, it wastes time measuring random, likely poor, schedules.
To manage this, auto-tuning processes usually employ an epsilon-greedy strategy or Upper Confidence Bound (UCB) approach. In the early rounds of tuning, the system randomly samples the search space to build a diverse dataset. As the cost model becomes more accurate (lower training error), the system shifts toward selecting the top-predicted candidates to converge on the optimal schedule.
Interestingly, a cost model does not need to predict the exact execution time perfectly to be useful. It only needs to rank the schedules correctly. If Schedule A is faster than Schedule B on hardware, the cost model must predict .
Because relative ordering is more important than absolute precision, many auto-tuners use ranking loss functions (like pairwise ranking loss) during training rather than standard Mean Squared Error (MSE). This ensures the model focuses on distinguishing good schedules from bad ones, rather than trying to predict the exact number of nanoseconds a kernel will take.
The following chart demonstrates how a cost model improves over tuning trials. Initially, the correlation between predicted rank and actual hardware rank is low. As more data is collected, the model aligns closer to reality.
Improvement in ranking correlation as the auto-tuner collects more hardware measurements.
When configuring a cost model for a specific target, different factors influence the model's complexity:
By effectively using a cost model, we decouple the search for efficiency from the physical limitations of benchmarking. This allows the compiler to evaluate billions of potential programs, eventually generating code that rivals or exceeds hand-tuned libraries like cuDNN or MKL.
Was this section helpful?
© 2026 ApX Machine LearningAI Ethics & Transparency•