机器学习成本模型

在物理硬件上评估数千个候选调度会在编译流程中造成明显瓶颈。虽然遗传算法或模拟退火求解器可以快速生成有效配置，但内核在GPU或TPU上的实际运行需要毫秒到秒级时间。当搜索空间包含数十亿种可能性时，硬件在环验证在计算上会变得难以负担。

为弥合庞大搜索空间与有限硬件可用性之间的差异，自动调优器使用统计成本模型。成本模型作为硬件的高吞吐量 (throughput)替代品，无需执行即可预测特定调度的性能。这使搜索算法每秒能查询数千个候选的性能，从而筛选出最有潜力的配置进行实际硬件基准测试。

替代模型的作用

正式来说，我们将优化问题定义为从有效调度空间 $S$ 中找到一个调度 $s$，使其在硬件目标 $H$ 上的运行时间最小。真正的成本函数 $Cost(s, H)$ 计算起来很昂贵。成本模型提供了一个近似函数 $\hat{f}(s)$，满足：

$\hat{f}(s) \approx Cost(s, H)$

搜索算法主要与 $\hat{f}(s)$ 进行交互。该工作流程采用主动学习方法。调优器最初不了解硬件对特定算子的响应。它随机选择一小批调度，在设备上进行测量，并使用得到的延迟数据来训练初始成本模型。

随着调优的进行，模型会预测下一批候选的性能。最佳候选在硬件上进行测量，这些真实测量数据被添加到训练集中以更新模型。这个反馈循环不断提高模型在搜索空间中关键区域的准确性。

自动调优中的主动学习循环。成本模型作为过滤器，使搜索策略能够查看数百万个候选，同时只在物理硬件上运行最有潜力的那些。

程序特征工程

机器学习 (machine learning)模型需要数值输入，但循环调度是一个结构化的语法对象。将抽象语法树 (AST) 或特定循环配置转换为特征向量 (vector)，是构建有效成本模型的重要一步。

统计特征提取

在像 AutoTVM 这样的系统中，特征提取依靠专家定义的规则。编译器遍历循环嵌套的低级中间表示 (IR)，并提取统计计数器。常见特征包括：

• 内存访问模式： 连续与步进式内存访问的数量。这有助于模型预测缓存一致性和内存带宽利用率。
• 算术强度： 浮点运算量与传输内存字节数之比。
• 循环结构： 循环嵌套的深度、每个循环的范围以及展开因子。
• 向量化 (quantization)： 向量指令的长度和 SIMD 通道的利用率。

这些特征被平展成固定长度的向量，并输入到梯度提升决策树 (GBDT) 中，例如 XGBoost。树模型在此处特别有用，因为它们处理特征之间的非线性关系（例如，瓦片大小与缓存大小），而无需数据归一化 (normalization)。

基于深度学习 (deep learning)的提取

Ansor (自动调度器) 等新型系统使用深度学习来自动学习特征表示。这些系统不是使用手工设计的计数器，而是将程序视为图或序列。

例如，调度可以被编码为一棵树，其中每个节点代表一个循环或语句。Tree-LSTM (长短期记忆) 网络或图神经网络 (neural network) (GNN) 处理这种结构以生成嵌入 (embedding)向量。这种嵌入捕捉到复杂的依赖关系，例如定义与其使用之间的距离，这在简单的统计计数器中可能会丢失。

排序而非回归

虽然目标是最小化延迟，但成本模型不一定需要精确预测以毫秒计的运行时间。更重要的是，模型能正确地对两个调度的相对性能进行排序。如果调度 A 在硬件上比调度 B 快，即使绝对值不准确，成本模型也应预测 A 的分数低于 B。

因此，现代成本模型通常优化成对排序损失，而非标准均方误差 (MSE) 回归损失。

给定两个调度 $s_i$ 和 $s_j$，其测量延迟分别为 $y_i$ 和 $y_j$，预测分数分别为 $\hat{y}_i$ 和 $\hat{y}_j$，RankNet 风格的损失函数 (loss function)将错误排序的概率降到最低。我们使用 sigmoid 函数 $\sigma$ 定义 $s_i$ 比 $s_j$ 快的概率为：

$P_{ij} = \frac{1}{1 + e^{-(\hat{y}_i - \hat{y}_j)}}$

损失函数会惩罚预测概率与硬件上观察到的实际排序之间的差异。这种方法使模型对硬件测量中的噪声具有抵抗力，并将学习能力集中在区分好坏调度上，而不是学习 GPU 的精确物理特性。

搜索空间效率

使用成本模型带来的效率提升非常显著。卷积神经网络 (neural network) (CNN) 层的典型调优会话可能涉及查看大小为 $10^{10}$ 的搜索空间。

没有成本模型，随机搜索或简单启发式方法可能会停滞不前，只找到局部最优解。有了成本模型，搜索过程能够有效地为每次实际硬件测量“模拟”数千步。下表展示了不同搜索策略的收敛行为。

调优收敛的比较。成本模型（蓝色）快速识别高性能区域，在 300 次试验内达到接近峰值的吞吐量 (throughput)，而随机搜索（灰色）改进缓慢。

成本模型中的迁移学习 (transfer learning)

标准成本模型的一个明显局限是它们通常是任务特定的。训练用于优化形状为 $224 \times 224$ 的 Conv2d 算子的模型，通常无法预测 MatMul 甚至不同形状 Conv2d 的性能。这要求自动调优器为神经网络 (neural network)中的每个新子图从头开始。

为解决此问题，高级编译器采用迁移学习。通过在包含各种算子和调度的大量数据集上进行训练，一个基础模型能学习到代码效率的通用原则，例如非对齐 (alignment)内存访问的惩罚或寄存器分块的好处。在调优新算子时，系统会使用这个预训练 (pre-training)模型开始，并用少量逻辑样本进行微调 (fine-tuning)。这显著减少了“冷启动”时间，使调优器能够用更少的硬件测量找到最优调度。