TVM作为一个端到端的深度学习编译器堆栈，其核心优势之一就是能够将高级算子描述编译到多种硬件后端上高效执行。然而，算子级的高性能实现往往需要精心设计调度（Schedule），例如循环分块、向量化、线程绑定、内存局部性优化等。手动为每一种算子、每一种硬件平台编写调度模板几乎是不可能完成的任务。Ansor作为TVM中的自动调优组件，改变了这一局面，它不再依赖预定义的模板库，而是通过搜索的方式自动生成调度空间，并结合机器学习成本模型筛选出性能优异的候选方案。下面这张图展示了Ansor整体工作流程中的关键阶段。
要理解Ansor的价值，可以先回顾一下传统自动调优的痛点。早期AutoTVM依赖用户手动编写调度模板，模板中定义了可调参数的范围，然后通过模拟退火或遗传算法在参数空间内搜索。这种方式虽然有效，但模板的编写需要深厚的领域知识，而且一套模板往往只能覆盖少数几种算子形状和硬件特性。一旦遇到新的算子或者新的硬件架构，模板就需要重新编写。Ansor则从更底层的循环结构出发，自动推导出合法的调度草图，并在草图基础上随机填充具体变换参数，从而构建出远超模板覆盖范围的搜索空间。
Ansor调度搜索的分层机制
Ansor的调度搜索可以拆解成三个层次：草图生成（Sketch Generation）、随机注释（Random Annotation）和进化搜索（Evolutionary Search）。草图生成阶段的目标是为给定的计算表达式推导出一组高层调度结构。这些草图决定了计算的基本并行策略和数据访问模式，例如是否将某个循环维度绑定到GPU线程块、是否对某个轴进行循环分块以便利用共享内存。草图本身并不包含具体的分块因子或线程数量，它只描述调度的骨架。
以矩阵乘法为例，计算表达式是一个三重循环，分别遍历M、N、K三个维度。草图生成器会分析这个表达式的数据复用关系，自动推导出多种可能的草图。一种常见的草图是：将M和N维度映射到GPU的线程块和线程，K维度作为归约轴保留在循环内部；另一种可能是将K维度分块后放入共享内存，以减少全局内存访问次数。草图生成器会基于TVM的Halide IR进行模式匹配，识别出归约轴、逐元素轴等结构，并应用一系列预定义的推导规则来产生候选草图。这些规则包括多级分块、向量化、绑定到并行虚拟线程等。
草图生成之后，随机注释阶段负责为草图中的每个可调参数填充具体数值。例如，对于“循环分块”这一结构，随机注释会从预设范围内随机选择分块因子，同时决定是否启用循环展开以及展开次数。为了保证搜索空间的多样性，Ansor会为同一个草图生成大量不同的随机注释版本。这些注释后得到的完整调度会被逐一编译，并在目标硬件上实际执行以获取真实性能数据。实际测量得到的运行时间将作为后续成本模型训练和进化搜索的标签。
当累积了足够多的实测样本后，进化搜索阶段开始发挥作用。进化算法维护一个调度种群，通过交叉和变异产生新个体。交叉操作可以交换两个调度的某些子结构，变异操作则对单个调度中的某个参数进行随机扰动。新生成的调度会先经过成本模型预测性能，只有预测结果足够好的个体才会被送到硬件上实测，从而减少昂贵的真实测量次数。成本模型通常采用梯度提升树或轻量级神经网络，输入特征包括循环结构的统计信息、内存访问模式、算术强度等。进化搜索迭代到一定代数后，种群中性能最佳的调度即作为该算子的最终选择。
硬件后端适配的关键点
Ansor生成调度时并非孤立的计算图变换，它必须与目标硬件的执行模型紧密配合。不同硬件后端对调度的约束差异巨大，比如GPU上需要显式管理线程层次和共享内存，CPU上则更关注SIMD向量化和缓存分块，而专用加速器可能要求固定大小的张量核心指令。Ansor在草图生成阶段就引入了与后端相关的推导规则。
对于CUDA后端，草图生成器会优先考虑将外层循环绑定到blockIdx和threadIdx，同时识别归约轴是否适合放入共享内存。例如在卷积算子的隐式GEMM调度中，Ansor会尝试将输入特征图和权重分块缓存在共享内存中，减少对全局内存的重复读取。此时草图中的“分块”和“绑定”操作会直接对应到CUDA编程模型中的blockDim、shared memory分配以及同步原语。随机注释生成的具体分块因子必须满足硬件限制，比如共享内存大小不能超过48KB或动态共享内存上限，线程块维度乘积不能超过1024等。Ansor内置了这些硬件约束检查，不符合约束的调度会被直接丢弃，避免无效的编译和运行。
在适配CPU后端时，Ansor更注重循环的向量化和缓存友好性。草图生成器会推导出适合AVX2或AVX512指令集的内层循环结构，随机注释则决定向量化宽度、循环展开因子以及多级分块的大小。例如对于矩阵乘法，CPU调优可能生成这样的结构：最外层对M和N进行多级分块，中间层对K进行分块以利用L2缓存，最内层则使用向量化指令一次处理8个浮点数。这些参数的具体值会通过实际运行来验证，因为缓存容量、内存带宽和指令延迟都会影响最优配置。
张量核心（Tensor Core）的适配是另一个典型场景。Ansor可以生成调用WMMA（Warp Matrix Multiply Accumulate）指令的调度，但要求矩阵分块尺寸必须是16x16x16或8x8x4等固定形状，并且数据在共享内存中的布局需要满足特定的跨步模式。草图生成器会识别出矩阵乘法模式，优先尝试生成符合张量核心约束的草图，例如将K维度以16为步长分块，将M和N维度以8为步长映射到warp内的线程。随机注释此时只需在有限的可选值中选择，因为硬件强制了这些形状。实测结果表明，启用张量核心的调度相比普通CUDA核心实现通常能获得2到5倍的吞吐量提升。
实战：使用Ansor进行矩阵乘法自动调优
下面通过一个完整的Python示例展示如何在TVM中使用Ansor对矩阵乘法算子进行自动调优。首先需要定义计算表达式和调优任务。代码中使用了TVM的Tensor Expression（TE）来描述矩阵乘法，然后创建AutoTVM调优任务，并指定目标硬件为NVIDIA GPU（CUDA后端）。
import tvm
from tvm import te, auto_scheduler
# 定义矩阵乘法计算
M = N = K = 1024
A = te.placeholder((M, K), name='A')
B = te.placeholder((K, N), name='B')
k = te.reduce_axis((0, K), name='k')
C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name='C')
# 创建调优任务
target = tvm.target.Target("cuda")
task = auto_scheduler.SearchTask(func=C, args=(A, B), target=target)
# 配置调优参数
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=200, # 实际测量次数
num_measures_per_round=64, # 每轮测量数量
verbose=2,
measure_callbacks=[auto_scheduler.RecordToFile("matmul_ansor.json")],
)
# 启动自动调优
task.tune(tune_option)
上述代码中，num_measure_trials控制实际在硬件上运行的调度总数，这个值越大，搜索到高性能调度的概率越高，但耗时会线性增加。num_measures_per_round则影响进化搜索的批量大小，较大的批量有利于并行测量。调优过程会持续打印当前搜索到的最佳性能，并将所有测量记录保存到JSON文件中。这些记录后续可以用于训练更精确的成本模型，也可以直接加载最佳调度进行部署。
调优完成后，可以通过以下代码获取最佳调度并构建可部署的运行时模块。如果不想每次重新调优，可以从已保存的JSON记录中恢复最佳配置，直接编译运行。
# 获取最佳调度并编译
sch, args = task.apply_best("matmul_ansor.json")
func = tvm.build(sch, args, target)
# 分配设备内存并运行
import numpy as np
dev = tvm.cuda(0)
a_np = np.random.uniform(size=(M, K)).astype("float32")
b_np = np.random.uniform(size=(K, N)).astype("float32")
c_np = np.zeros((M, N)).astype("float32")
a_tvm = tvm.nd.array(a_np, dev)
b_tvm = tvm.nd.array(b_np, dev)
c_tvm = tvm.nd.array(c_np, dev)
func(a_tvm, b_tvm, c_tvm)
tvm.testing.assert_allclose(c_tvm.numpy(), np.dot(a_np, b_np), rtol=1e-3)
在这个例子中，Ansor会自动推导出适合CUDA的调度草图，包括将M和N维度映射到线程块和线程、将K维度分块放入共享内存。如果目标硬件是CPU，只需将target改为"llvm -mcpu=skylake-avx512"，Ansor就会生成完全不同的调度结构，重点优化向量化和多级缓存分块。这种后端自适应的能力正是Ansor相比固定模板调优的主要优势。
调优过程中的常见问题与优化策略
虽然Ansor大幅简化了自动调优流程，但在实际使用中仍有一些需要特别注意的地方。首先是搜索时间的控制。对于复杂算子或大规模输入形状，调优空间可能非常庞大，200次测量也许远远不够。此时可以通过设置更长的调优时间，或者采用基于历史数据的成本模型预热来加速收敛。TVM提供了auto_scheduler.LocalRunner和auto_scheduler.RPCRunner，可以并行使用多个GPU设备进行测量，显著缩短总体调优时间。
另一个常见问题是调度在特定输入形状上过拟合。Ansor调优通常针对固定的输入张量形状进行，如果生产环境中的形状会动态变化，需要针对多个典型形状分别调优，或者使用动态形状的搜索任务。TVM的auto_scheduler支持定义带有动态维度的任务，搜索空间会考虑符号化维度的影响，但代价是搜索复杂度上升。实际工程中常用的做法是针对最常见的几个形状分别调优，然后在运行时根据输入形状选择对应的预编译内核。
此外，成本模型的准确性直接影响进化搜索的效率。TVM内置的成本模型是梯度提升树，它在样本量较小时容易给出过于乐观的预测，导致许多实际表现不佳的调度被选中进行昂贵测量。可以通过增加初始随机测量的数量、或者使用自定义的性能特征（比如从调度IR中提取的循环嵌套统计信息）来改善。在一些自研硬件上，还可以结合硬件模拟器的反馈来替代真实测量，但这需要额外的集成工作。总体而言，Ansor的自动调优框架为不同硬件后端的算子优化提供了一条可扩展的路径，理解其搜索机制和适配约束，能够帮助开发者在实际项目中更高效地获得高性能算子实现。