Running an Auto-Tuning Session

Effective optimization relies on establishing a search space with defined parameters and a well-defined cost model. While these elements provide the foundational infrastructure, the actual execution of the search—known as an auto-tuning session—is where performance gains are ultimately achieved. An auto-tuning session acts as an experimental loop. The compiler iteratively generates candidate kernels, measures their performance on the target hardware, and uses that feedback to refine its search strategy.

The Architecture of a Tuning Session

A tuning session is not a monolithic process. It involves the coordination of several distinct components: the task scheduler, the code builder, the runner, and the measure callback. Understanding how these pieces interact is necessary for debugging performance issues or configuring the tuner for novel hardware.

In a typical setup, the machine running the compiler (the host) may be different from the machine executing the code (the target). This is common in edge AI, where you compile on a high-power workstation but tune for a mobile phone or an embedded accelerator. The tuner relies on an RPC (Remote Procedure Call) mechanism to send compiled artifacts to the device and receive timing data back.

The following diagram illustrates the cyclical data flow during an auto-tuning session. The optimizer generates configurations, which are compiled and sent to the hardware runner. Execution metrics are then fed back to update the cost model.

Extracting Tuning Tasks

Deep learning models consist of dozens or hundreds of operators. However, not every operator requires tuning. Many operations, such as Reshape or ReLU, are memory-bound and have limited optimization headroom compared to compute-bound operations like Conv2d or MatMul.

The first step in a session is task extraction. The compiler traverses the computation graph and identifies subgraphs that match specific patterns, usually heavy arithmetic operators coupled with their surrounding element-wise operations (due to operator fusion). Each unique subgraph structure becomes a "tuning task." If a model uses the same convolution layer structure ten times, the tuner only needs to optimize that kernel once and reuse the schedule.

We typically define a tuning configuration object that specifies the target hardware constraints. For example, when tuning for a GPU, the task must be aware of the maximum number of threads per block and shared memory limits.

configuring the Search Policy

Once tasks are identified, you must select a search policy. This determines how the tuner navigates the search space defined in the previous section.

1. Grid Search: Exhaustively tries every valid configuration in the search space. This is generally impractical for modern neural networks due to the combinatorial explosion of tiling and vectorization parameters.
2. Random Search: Samples configurations uniformly from the space. This acts as a baseline but often fails to find peak performance in reasonable timeframes.
3. Model-Based Search (Bayesian/XGBoost): The industry standard for production. The tuner trains a surrogate model (the cost model) to predict the rank of candidate schedules. It balances exploration (trying parts of the space with high uncertainty) and exploitation (refining variations of the current best candidates).

A typical configuration in Python pseudo-code might look like this:

tuning_option = {
    "n_trials": 2000,          # Maximum number of configurations to test
    "early_stopping": 600,     # Stop if no improvement after N trials
    "measure_option": {
        "builder": LocalBuilder(),
        "runner": RPCRunner(
            key="android_device",
            host="0.0.0.0", 
            port=9190,
            number=5,         # Average 5 runs per config
            repeat=1          # Measure the loop 1 time
        ),
    },
    "tuner": "xgb",           # Use XGBoost-based tuner
}

executing the Tuning Loop

When the tune command is invoked, the loop begins. The console output of an auto-tuner is often verbose, streaming real-time metrics. It is important to interpret this data correctly.

You will typically see a metric such as GFLOPS (Giga Floating Point Operations Per Second) or latency (in milliseconds). The goal is to maximize GFLOPS or minimize latency. Early in the session, performance is usually low as the tuner randomly samples the space to initialize the cost model. As the cost model acquires data (training samples), the tuner begins to suggest better candidates, and you should observe a steep upward trend in performance.

Eventually, the curve plateaus. This indicates that the tuner has likely found the hardware limit for that specific operator or is stuck in a local optimum.

The chart below demonstrates a typical convergence curve for a matrix multiplication operator. The blue dots represent individual trials, while the orange line tracks the best performance found so far. Note the rapid improvement in the first 50 trials as the cost model learns.

Handling Tuning Artifacts

The immediate output of a tuning session is not the compiled binary, but a tuning log. This is typically a JSON or text file containing the history of every configuration tried and its corresponding execution time.

This log file is a valuable asset. It effectively serves as a database mapping operator signatures (input shapes and data types) to their optimal loop schedule parameters.

1. Inspection: You can manually inspect the top-performing configuration to understand what the compiler "learned." For instance, you might find that for a specific GPU, the tuner consistently chooses a tile size of $16 \times 16$ rather than $32 \times 32$ to avoid register spilling.
2. Transfer Learning: If you have a tuning log from a similar hardware device, you can use it to seed the cost model for a new session. This drastically reduces the time required to find good schedules on the new device.

Applying the Best Schedule

The final phase of the session is the compilation of the optimized model. This differs from the standard compilation flow because the compiler now has an external source of truth.

When you invoke the high-level build function (e.g., relay.build in TVM or torch.compile), you provide the tuning log as a context manager or argument. As the compiler lowers the computation graph, it queries the log for each operator. If a matching entry is found, the compiler bypasses its default heuristics and applies the specific loop transformations, tiling factors, vectorization lengths, and unrolling factors, defined in the optimal configuration.

If no entry is found for a specific operator (perhaps the tuning session was interrupted or the operator was excluded), the compiler falls back to a default schedule. This ensures the model remains functional, though those specific layers will not benefit from the optimization boost.

The resulting machine code is a hybrid of highly tuned kernels for the computationally intensive parts of the graph and generic implementations for the rest. This optimized binary is then ready for deployment, offering lower latency and better energy efficiency than the unoptimized baseline.