搜索算法：从随机到遗传

找到循环嵌套的最佳调度与训练神经网络 (neural network)的权重 (weight)有着根本的不同。在模型训练中，优化面是连续且可微分的，这使得梯度可以引导优化器趋向最小值。在编译器优化中，搜索空间是离散的、不可微分的，并且常常是不连续的。参数 (parameter)的一个小改动，例如将瓦片大小增加一，可能导致缓存冲突，从而使性能急剧下降；而稍大一点的改动则可能与缓存行完美对齐 (alignment)，使吞吐量 (throughput)翻倍。

由于我们无法计算相对于调度参数（如 tile_x 或 unroll_factor）的梯度，我们依赖于黑盒 (black box)优化算法。这些算法从上一节定义的搜索空间中抽取配置样本，并通过实际硬件测量或成本模型来评估其质量。本节审视从简单的随机方法到现代自动调度器（如Ansor）中使用的复杂演化策略的进展。

基准：随机搜索

随机搜索作为评估任何自动调优系统有效性的标准基准。它通过从定义的搜索空间中均匀地抽取有效配置样本来运行。随机搜索虽然看起来原始，但它具有使其具有应用价值的独特属性：

1. 无偏寻优： 它不会陷入局部最优，因为它没有先前状态的记忆。
2. 并行性： 每个样本都是独立的，允许在分布式工作器之间进行大规模并行化。
3. 搜索空间覆盖： 它提供了整个空间中性能分布的统计近似。

在高维空间 (high-dimensional space)中，随机搜索常常优于网格搜索，因为网格搜索会浪费资源在那些可能不会显著影响目标函数的维度上进行寻优。然而，对于编译器调度而言，随机搜索效率低下。调度空间的高性能区域通常极其狭窄，通常只有不到0.01%的有效调度能达到峰值性能的90%以内。单纯依靠概率进入这个“好”区域的计算成本很高。

模拟退火

为了改进随机采样，我们需要一种方法，它能对空间进行考察，并最终收敛到高性能区域。模拟退火（SA）是一种概率技术，能够摆脱局部最优，这是循环调度这种复杂优化方面常见的问题。

SA维护一个当前状态 $s$，并迭代地提出一个相邻状态 $s'$。如果 $s'$ 的性能更好，它将被接受。如果 $s'$ 的性能更差，它将以玻尔兹曼分布定义的概率 $P$ 被接受：

$P( ext{接受}) = \exp\left( \frac{E(s) - E(s')}{T} \right)$

在这里，$E(s)$ 是能量（性能的倒数），$T$ 是温度。温度开始时很高，这使得算法可以接受较差的配置并广泛地寻优空间。随着搜索的进行，$T$ 会根据降温方案衰减。当 $T$ 趋近于零时，算法的行为就像一个贪婪的爬山算法，只接受改进。

在TVM或MLIR的背景下，“相邻状态”是通过变异当前调度的一个或多个特定参数 (parameter)来定义的，例如，将向量 (vector)化因子从4改为8，或交换两个循环的顺序。

演化算法和遗传算法

现代深度学习 (deep learning)编译器，包括Ansor (AutoTIR) 系统，大量依赖演化算法。与模拟退火只跟踪单个状态不同，演化算法维护一个候选调度种群。这种方法对张量程序尤其有效，因为有效的优化常常需要特定特性的组合（例如，特定的瓦片结构与特定的向量 (vector)化策略结合），这些特性可以独立演化然后组合。

编译器自动调优中的演化过程通常遵循这些步骤：

1. 初始化： 种群会被初始化，通常混合了随机有效调度和基于启发式的调度。
2. 选择： 最适应的个体（预测延迟最低的那些）被选择来繁育下一代。
3. 交叉： 两个父调度中的属性被组合。对于循环嵌套，这可能意味着取父代A的瓦片结构和父代B的线程绑定策略。
4. 变异： 应用随机变动以保持遗传多样性。这可以防止种群过早地收敛到次优的局部峰值。

自动调度中演化搜索策略的工作流程。种群通过变异和交叉进行演化，并由成本模型或硬件执行得出的适应度函数引导。

代码生成中的变异操作符

在代码生成中实现变异比一般优化要求更高，因为生成的调度必须保持合法。无效的变异可能破坏数据依赖或违反硬件约束（例如使用超出可用量的共享内存）。

常见的变异操作符包括：

• 瓦片大小变异： 将瓦片大小乘以或除以一个因子（例如，将块大小从32改为64）。
• 重排变异： 随机交换循环嵌套中两个相邻的循环。
• 标注变异： 改变诸如 unroll 深度或 vectorize 宽度等编译指示。
• 并行化变异： 改变循环与GPU线程（blockIdx/threadIdx）的绑定。

执行交叉时，编译器必须确保组合特性不会冲突。例如，如果父代A对归约轴进行瓦片化而父代B没有，盲目合并它们可能导致IR畸形。

寻优与优化现有方案

设计这些搜索算法的核心难题是平衡寻优（找到新的、可能更好的吸引盆地）和优化现有方案（细化目前找到的最佳解决方案）。

• 随机搜索 是100%寻优，0%优化现有方案。
• 贪婪搜索 几乎是0%寻优，100%优化现有方案。
• 演化算法 提供可调的平衡。变异率控制寻优，而选择压力控制优化现有方案。

实践中，自动调度器常采用多级方法。它们可能从高寻优阶段（如随机搜索或高温SA）开始，以用多样化的数据点为成本模型提供初始数据。一旦成本模型变得相当准确，系统就会切换到演化策略，以积极优化最佳候选。

以下图表呈现了这些算法在GPU上调优矩阵乘法核时典型的收敛行为。

吞吐量 (throughput)随时间收敛的比较。随机搜索难以从基准提高。模拟退火分步改进，但可能会停滞。演化搜索通过结合高性能候选的最佳特性而持续攀升。

从搜索到成本模型

尽管演化算法很强大，但它们需要数千次评估才能收敛。在GPU上编译和运行数千个核速度很慢，单个操作符可能需要几分钟甚至几小时。这种延迟对于即时 (JIT) 编译场景是不可接受的。

为了解决这个问题，现代框架不会执行演化搜索生成的每个候选。相反，它们查询统计成本模型。搜索算法生成一个候选，成本模型预测其性能，只有最有希望的候选最终才会在硬件上编译和测量，以微调 (fine-tuning)模型。搜索算法与成本模型之间的这种互动使得编译器能够高效地遍历指数级优化空间。