TVM Auto-Scheduler (Ansor)

The core of Ansor is a learned cost model that predicts the execution latency of a candidate program without running it on hardware. This model is trained on measured performance data from sampled programs. The search algorithm uses this model to prioritize exploration of the most promising regions of the schedule space, which is astronomically large. This approach is far more efficient than exhaustive search or template-based autotuning.

• The cost model is typically a gradient-boosted tree (e.g., XGBoost) or a neural network.
• It evaluates features extracted from the program's loop structure, memory access patterns, and hardware characteristics.
• This guidance allows Ansor to discover novel, high-performance schedules that human experts might overlook.

Ansor decomposes the optimization problem using a hierarchical representation of the computational graph. It does not schedule low-level loops directly from the entire model. Instead, it:

1. Sketches High-Level Structures: Generates possible high-level loop structures (sketches) for each computational task (operator), such as whether to tile, vectorize, or unroll loops.
2. Explores Low-Level Parameters: For each sketch, it searches for optimal low-level parameters like tile sizes, unroll factors, and parallelization strategies.

This hierarchy makes the search tractable by first deciding on major optimization decisions before fine-tuning the numerical parameters.

To train its cost model, Ansor must generate a diverse set of training programs. It uses automated program sampling to create this dataset.

• Random Annotation: It starts by randomly applying a set of schedule primitives (like split, reorder, tile) to the initial computational graph.
• Evolutionary Search: It then uses an evolutionary algorithm to mutate and recombine high-performing sampled programs to generate new, potentially better candidates.
• This process continuously generates new programs, measures their actual hardware performance, and uses the results to refine the cost model in an iterative feedback loop.

Ansor's search is fundamentally guided by the target hardware's constraints and capabilities. It automatically discovers schedules optimized for specific memory hierarchies, cache sizes, vector units, and parallel compute units.

• It learns hardware-specific patterns, such as optimal tile sizes for different levels of cache (L1, L2, shared memory).
• It automatically applies hardware-specific optimizations like vectorization for SIMD units and thread binding for multi-core CPUs or GPU streaming multiprocessors.
• This eliminates the need for manual, platform-specific tuning, enabling a single model description to achieve near-peak performance across diverse backends (e.g., ARM CPU, NVIDIA GPU, x86).

Ansor is not a standalone tool; it is deeply integrated into the TVM compiler stack. It operates on TVM's intermediate representation (IR), specifically the Tensor Expression (TE) language.

• Input: Takes a computational graph defined via TVM's Relay frontend or directly as TE expressions.
• Output: Produces an optimized schedule object, which is then lowered by TVM to generate hardware-specific code (LLVM, CUDA, OpenCL, etc.).
• This integration allows Ansor to be used as a drop-in replacement for manual scheduling within existing TVM-based deployment pipelines, providing automatic performance gains.

Ansor was developed to address limitations in TVM's earlier autotuning system, AutoTVM. Key differentiators are:

• Template-Free: AutoTVM required manual definition of schedule templates for each operator, which limited the search space. Ansor requires no templates; it explores the space of all valid schedules.
• End-to-End Graph Optimization: AutoTVM tuned operators in isolation. Ansor performs cross-operator optimizations, like fusion, across the entire computational graph during its search.
• Superior Search Efficiency: The learned cost model and hierarchical search in Ansor typically find better-performing schedules faster than AutoTVM's template-guided search and simpler cost models.

Researchers use Ansor as a performance evaluation tool for new AI accelerators and exotic architectures. Its ability to automatically explore the schedule space eliminates the need for a mature manual kernel library, allowing for:

• Rapid architectural exploration by measuring the performance potential of new tensor instructions or memory hierarchies.
• Benchmarking novel operators (e.g., from cutting-edge research papers) without months of manual optimization effort.
• Serving as the default optimization backend in ML compilers for academic hardware simulators and FPGA prototypes.

Ansor excels where no pre-optimized kernels exist, making it indispensable for newly published model architectures. Examples include:

• Vision Transformers (ViTs) and their variants with custom attention mechanisms.
• Diffusion models with unique sampling schedules and operator compositions.
• Graph Neural Networks (GNNs) with sparse, irregular computation patterns that are poorly served by standard GEMM libraries.
• Custom, domain-specific layers developed in-house by AI research teams.

Ansor acts as the low-level optimization engine within larger Neural Architecture Search (NAS) or hyperparameter tuning frameworks. In this role:

• A Hardware-Aware NAS algorithm proposes a candidate model architecture.
• Ansor is invoked to rapidly estimate the latency of that model on the target hardware by generating and benchmarking a schedule.
• This feedback loop allows NAS to directly optimize for inference speed or energy consumption, not just accuracy. Systems like Once-For-All (OFA) can use Ansor to profile the performance of different sub-networks.

Neural Architecture Search (NAS) is an automated process for designing optimal neural network architectures for a specific task and hardware constraint. It uses search algorithms (e.g., reinforcement learning, evolutionary algorithms) to explore a vast space of possible model configurations (layer types, connections, channel widths). While Ansor searches for optimal schedules (how to execute a fixed graph), NAS searches for the optimal graph structure itself. Hardware-Aware NAS directly incorporates metrics like latency or power consumption into the search objective.

Graph Compilation is the overarching process of transforming a high-level neural network computational graph (e.g., from PyTorch or TensorFlow) into an optimized, hardware-specific sequence of low-level operations. This involves several stages where Ansor operates:

• Operator Lowering: Decomposing framework-specific layers into primitive tensor operations (e.g., conv2d -> im2col + GEMM).
• Graph-Level Optimizations: Applying transformations like operator fusion, constant folding, and dead code elimination.
• Schedule Generation & Auto-Tuning: This is Ansor's primary domain—finding the optimal loop nest structure, parallelization, and memory layout for each operator.

Kernel Auto-Tuning is the automated, empirical search for the optimal configuration parameters of a single computational kernel. This is a critical sub-problem that Ansor solves at scale. For each candidate schedule, Ansor must generate code and find the best:

• Thread block size and grid dimensions for GPU targets.
• Loop tiling factors for CPU and NPU cache hierarchies.
• Vectorization and unrolling parameters. Ansor automates this by generating many code variants for each schedule and using a cost model to predict performance, followed by actual measurement on hardware to refine the model.

Operator Fusion is a fundamental compiler optimization that combines multiple sequential neural network operations into a single, fused kernel. This is a key transformation applied during graph compilation, often before Ansor's scheduling phase. Examples include:

• Fusing a convolution, bias addition, and ReLU activation.
• Fusing a matrix multiplication with a subsequent non-linearity. Benefits:
• Minimizes intermediate memory writes/reads between kernels, reducing bandwidth pressure.
• Reduces kernel launch overhead by executing fused logic in a single launch.
• Increases operational intensity, making the kernel more compute-bound.

These are two deployment strategies for compiled models, both supported by TVM and Ansor.

Just-In-Time (JIT) Compilation: Code is compiled at runtime, immediately before execution. This allows for optimizations specific to dynamic input shapes or target hardware discovered at deployment. Ansor can run its search during JIT compilation.

Ahead-Of-Time (AOT) Compilation: Code is fully compiled into a standalone, minimal-footprint executable binary before runtime. This is critical for edge deployment on resource-constrained devices. Ansor's search is performed during the AOT compilation phase, and the optimal schedule is baked into the final binary.

The Roofline Model is an analytical performance model used to guide and evaluate optimizations like those Ansor produces. It plots attainable performance (FLOPS/sec) against Operational Intensity (Ops/Byte).

Operational Intensity measures the arithmetic work done per byte of data transferred from main memory. Ansor's goal is to transform schedules to maximize this metric, pushing kernels toward the hardware's compute-bound roof. Techniques include:

• Loop tiling to increase data reuse from cache.
• Operator fusion to avoid writing intermediate results to memory.
• Selecting schedules that maximize cache locality.