自动调优中的成本模型

循环优化的搜索空间很大，仅一个矩阵乘法运算符，考虑到平铺大小、向量化因子和循环重排组合，就可能拥有数十亿个有效的调度配置。如果在实际硬件上测试单个配置即使只花费一小部分秒，评估整个空间在计算上也是不可行的。

为了有效地处理如此多的可能性，我们不能依靠硬件测量每个候选项。相反，我们使用成本模型。成本模型充当硬件的数学替代。它可以在不实际运行代码的情况下预测特定循环调度的性能。这使得自动调优器每秒可以模拟数百万种配置，筛选掉不佳的候选项，并找出有希望的候选项进行实际验证。

成本模型的作用

成本模型的基本目标是估算所生成内核的执行时间。如果我们将调度配置定义为 $S$，实际硬件执行时间为 $T(S)$，则成本模型 $f$ 旨在使误差最小化：

$f(S) \approx T(S)$

在典型的自动调优流程中，搜索算法提出一批候选调度。成本模型评估这些候选项，并为每个候选项分配一个分数或预测延迟。只有预测分数最佳的候选项才会被发送至硬件进行性能测试。这通过在过程早期剔除低效调度，显著缩短了调优时间。

机器学习编译器中使用的成本模型主要分为两类：分析模型和学习模型。

分析成本模型

早期的编译器大量依赖分析建模。这些模型使用基于硬件规格的预设公式。分析模型通过计算操作数量和基于指令成本估算延迟来工作。

例如，一个简单的CPU分析模型可能这样计算成本：

$\text{成本} = (N_{浮点运算} \times C_{浮点运算}) + (N_{加载} \times C_{内存}) + (N_{缺失} \times C_{惩罚})$

其中 $C$ 分别表示浮点运算、内存加载和缓存未命中所需的时钟周期成本。

虽然分析模型评估速度非常快，但它们往往无法描述现代加速器的复杂行为。现代 GPU 和 TPU 具有乱序执行、内存bank冲突和指令流水线等复杂行为，这些难以用静态公式表达。因此，分析上看起来高效的调度，在实际中可能表现不佳，这是由于共享内存bank冲突等硬件特定的瓶颈造成的。

学习成本模型

TVM 和 Halide 等现代框架已转向学习成本模型。这些系统不再依赖固定公式，而是使用机器学习来预测其他机器学习工作负载的性能。这种方法将性能预测视为回归或排序问题。

学习成本模型通常采用轻量级统计方法，例如 XGBoost（梯度提升树）或小型多层感知器 (MLP)。工作流程如下：

1. 特征提取： 编译器从循环调度中提取特征。这些特征可能包括内存访问次数、循环嵌套深度、向量化宽度和数据复用模式。
2. 训练： 模型在 $(调度, 执行时间)$ 对上进行训练。
3. 推理： 模型为新的、未见过的调度预测分数。

这种方法很有效，因为它能适应特定硬件。如果你在 NVIDIA A100 上调优模型，成本模型会学会 A100 的性能特征。如果你切换到移动 ARM CPU，成本模型会重新训练，学会该架构的独特特征。

循环调度的特征提取

为了训练模型，循环嵌套的抽象语法树 (AST) 必须转换成特征向量。ML 编译器通常提取的特征包括：

• 内存访问模式： 内存访问是连续的（有利于向量化）还是跨步的（导致缓存未命中）？
• 算术强度： 计算操作与内存操作的比率。
• 循环结构： 循环的大小和展开程度。
• 并行度： 使用了多少线程以及工作负载如何分配。

下图展示了数据流从搜索空间经过成本模型，再进入硬件测量阶段。

学习成本模型的迭代工作流程，展示了硬件测量如何随时间改进预测。

发掘未知与利用已知方案的权衡

自动调优中一个重要问题是发掘未知与利用已知方案之间的平衡。

• 利用已知： 选择成本模型预测速度快的调度。这依赖于模型当前已有的知识来寻找最佳性能。
• 发掘未知： 选择成本模型不确定的调度。这有助于模型了解其尚未遇到的搜索空间部分。

如果调优器只利用已知方案，它可能陷入局部最小值，因为成本模型在未发掘的区域不准确。如果它只发掘未知方案，就会浪费时间测量随机的、可能表现不佳的调度。

为此，自动调优过程通常采用 epsilon-greedy 策略或上限置信区间 (UCB) 方法。在调优的早期阶段，系统随机采样搜索空间以构建多样化数据集。随着成本模型变得更准确（训练误差更低），系统转向选择预测分数最高的候选项，以收敛到最优调度。

衡量指标：排名对比绝对延迟

有意思的是，成本模型不必完美预测精确的执行时间就能发挥作用。它只需要正确地对调度进行排序。如果调度 A 在硬件上比调度 B 快，成本模型必须预测 $f(A) < f(B)$。

由于相对顺序比绝对精度更重要，许多自动调优器在训练期间使用排序损失函数（如成对排序损失），而不是标准均方误差 (MSE)。这确保模型侧重于区分好坏调度，而不是试图预测内核将花费的精确纳秒数。

下图展示了成本模型在调优尝试中如何改进。最初，预测排名与实际硬件排名之间的相关性较低。随着收集更多数据，模型更接近真实情况。

排名相关性的提升，随着自动调优器收集更多硬件测量数据。

硬件注意事项

为特定目标配置成本模型时，不同的因素会影响模型的复杂程度：

1. CPU： 性能通常主要受缓存局部性和向量化影响。CPU的成本模型非常关注循环平铺大小及其与L1/L2缓存大小的关系。
2. GPU： 并行度和内存延迟隐藏很重要。成本模型必须学会偏好能很好地映射到线程块和warps的调度，同时最小化全局内存往返次数。
3. 边缘加速器： 这些设备通常具有有限的指令集或专用内存缓冲区。这里的成本模型依赖于严格遵守有效的指令约束，经常惩罚任何将数据溢出到较慢系统RAM的调度。

通过有效使用成本模型，我们将寻求效率的过程与性能测试的物理限制分离。这使得编译器能够评估数十亿个潜在程序，最终生成可与 cuDNN 或 MKL 等手调库媲美或超越的代码。