自动化调度搜索

定义搜索空间限定了优化的范围，但并未指明编译器应选择哪种配置。一个典型循环嵌套中的有效组合数量之多，常达数百万种以上。如果编译器要在实际硬件上编译并测试每一个变体，单个矩阵乘法核可能需要数周才能完成优化。为使此过程可行，机器学习编译器需借助高效的搜索算法。

这些算法在搜索空间中寻优，以便在合理的时间预算内找到高性能调度方案。目的是最大化一个目标函数，通常是吞吐量（FLOPS）或最小化延迟，通过迭代选择预计表现优于当前最佳的候选方案。

组合难题

思考一个简单的矩阵乘法操作，其中包含三个循环。为优化此操作，我们可能会应用循环分块、向量化和循环展开。如果我们把每个循环分成三级分块，且每个分块有10个候选大小，组合数量将呈指数级增长。

从数学角度看，如果我们有$k$个可调节参数$p_1, p_2, \dots, p_k$，并且每个参数有 $|D_i|$ 个有效选项，则搜索空间 $S$ 的大小为：

$|S| = \prod_{i=1}^k |D_i|$

即使是中等规模的循环嵌套，$|S|$ 也轻易超出 $10^9$。由于在硬件上测量单个配置需要可观的时间（编译 + 执行 + 开销），我们需要一种策略来找到全局最优，而不必遍历 $S$ 中的每个点。

试探与利用

编译器中的搜索算法通常平衡两个相互竞争的目标：

1. 试探： 探查搜索空间中的未知区域，以避免陷入局部最优。
2. 利用： 围绕已知的高性能配置精炼搜索，以榨取最后的性能潜力。

一种简单的方法是随机搜索，编译器从 $S$ 中随机均匀地抽取配置。虽然实现和并行化简单，但它忽略了问题的结构，不从先前的失败或成功中学习。因此，现代编译器很少完全依赖随机搜索。相反，它们使用进化算法或基于模型的方法。

进化搜索算法

进化算法（EA）将搜索过程视为生物进化。一个“种群”包含一组循环调度方案。在每次迭代（代）中，算法选择表现最佳的调度方案，并应用遗传算子来生成新的候选方案。

变异与交叉

在循环优化中，变异涉及随机更改调度方案的一个参数。例如，将分块大小从32更改为64，或更改循环展开因子。

交叉结合了两个高性能父代调度方案的特点。如果父代A在内存分块方面表现出色但向量化不佳，而父代B分块不佳但向量化出色，交叉操作可能会产生一个子代调度方案，继承A的分块和B的向量化特点。

经过多代演化，种群会向搜索空间中更高性能的区域演进。这种方法很有效，原因是有效的调度方案通常聚集在一起；如果分块大小为32效果良好，那么分块大小为16或64很可能比像7这样的随机值更好。

基于成本模型的搜索

最先进的自动调优系统，如TVM或Halide中的系统，采用基于模型的方法。此方法使用上一节中讨论的统计成本模型来引导搜索。编译器不是在硬件上运行每个候选方案，而是使用轻量级机器学习模型（成本模型）来预测数千个候选方案的性能。

工作流程分轮进行：

1. 候选方案生成： 搜索策略提出大量潜在的调度方案。
2. 预测： 成本模型预测此批次的运行时间。这很快，因为它完全在软件中完成。
3. 选择： 系统选择预测得分最高的 $k$ 个候选方案。
4. 测量： 这些被选中的少数方案在实际硬件设备上编译并运行，以获取真实执行时间。
5. 反馈： 测得的数据（调度特征 + 实际时间）被添加到训练数据集中。
6. 更新： 成本模型使用新数据重新训练或更新，以提高其下一轮的准确性。

基于模型的优化循环使得编译器在调优过程中得以学习硬件的性能特征。

这种方法样本效率远高于随机搜索。成本模型能迅速学会拒绝不良配置（例如导致缓存颠簸的分块大小），从而使实际基准测试时间仅用于有潜力的候选方案。

模拟退火与XGBoost

为高效地搜索空间，编译器常将全局搜索策略与局部策略结合。一种常见实现是使用模拟退火来搜索成本模型所界定的空间。成本函数作为能量模型，算法试图找到最低点（最小延迟）。

同时，梯度提升树（如XGBoost）常被用作成本模型本身。它们对编译器参数的离散和非线性特点很有效。搜索算法生成数千个候选方案，XGBoost模型对它们进行排名，只有排名前百分之几的会被发送到硬件执行器。

收敛与停止条件

自动调优会话无法无限期运行。我们必须定义停止条件，以判断搜索何时足够。常见条件有：

• 试验次数： 在硬件上测量 $N$ 个配置后停止（例如1000次试验）。
• 时间预算： 调优会话运行特定时长后停止（例如30分钟）。
• 提前停止： 如果性能在连续多次迭代后没有改善则停止。

下方的图表说明了不同搜索策略如何随时间收敛。请注意，与随机搜索相比，基于模型的搜索（引导式）通常在更少的硬件试验次数下获得更高的性能。

连续试验中的性能提升。引导式搜索使用历史数据更快地找到最佳参数。

搜索中的迁移学习

自动调优面临的一个难题是，它通常对每个新模型或算子都从头开始。为缓解此问题，现代研究集中于迁移学习。如果编译器已调优过 $1024 \times 1024$ 大小的矩阵乘法，那么在调优 $1024 \times 512$ 大小的乘法时，它应使用该知识。

在此情况下，搜索算法使用在类似工作负载上预训练过的权重来初始化成本模型。这使搜索获得“热启动”，减少了找到良好调度方案所需的时间。搜索不再盲目地遍历整个空间，而是立即关注过去对类似张量形状有利的区域。

实际实现

在TVM或MLIR等框架中配置自动化搜索时，您实质上是在配置这些算法的参数。您定义了特定的硬件目标（它决定了指令集和内存限制）、试验次数以及搜索策略的复杂程度。

通过理解这些搜索算法的机制，您可以更好地诊断调优问题。如果性能过早停滞，搜索可能陷入局部最优（过度利用）。如果性能不稳定且提升缓慢，搜索空间可能过大或成本模型不准确（无指导的过度试探）。优化优化器通常是部署高性能机器学习模型的最后一步。