定义搜索空间

定义搜索空间是自动化编译器优化的基础步骤。在优化器找出最佳调度前，它需明白有效配置的边界。搜索空间表示特定硬件目标上，给定算子的全部可行程序转换。如果空间过小，优化器可能错失最有效的优化。如果空间过大，搜索将变得计算成本高昂，且可能无法在合理时间内得出结果。

搜索空间的结构

在中间表示层，计算由数据数组上的嵌套循环定义。搜索空间包含可调参数，它们在不改变数值输出的前提下，调整这些循环的执行方式。这些参数通常分为三类：平铺因子、执行顺序和硬件绑定。

考虑一个简单的向量加法，我们遍历 $N$ 个元素。基线实现是一个单一的串行循环。为优化它，我们可应用循环平铺（将循环拆分成更小的块），以将数据放入 CPU 缓存，或将工作映射到 GPU 线程。

如何拆分循环的决定构成了我们搜索空间的第一维。如果 $N = 1024$，有效的平铺大小包括 1, 2, 4, ..., 512, 1024。如果我们把循环拆分成两层（外层和内层），关系定义为：

$i = i_{outer} \times \text{tile\_size} + i_{inner}$

这里，tile_size 是我们搜索空间中的一个变量。编译器必须为该变量定义一个有效值范围。在简单的整数网格搜索中，这一单个决定的空间包含 1024 的所有因子。

优化选项的构成

实际算子，如矩阵乘法或二维卷积，涉及多个嵌套循环。随着我们增加更多维度和优化原语，搜索空间会组合式增长。

对于涉及维度 $M, N, K$ 的矩阵乘法操作，编译器必须同时决定所有三个维度的平铺因子。此外，这些循环的嵌套顺序会影响内存局部性。这便将“重排序”引入了搜索空间。

下图展示了单个算子如何形成一个决策分支树，并在叶节点产生一个特定配置（或“候选”）。

优化搜索空间的结构，展示了形成有效调度所需的决策顺序。

约束和有效性

并非所有参数组合都能产生有效或高效的代码。定义搜索空间也涉及强制执行约束，以立即剔除无效候选。这能防止自动调优器浪费周期评估损坏的代码。

硬性约束

硬性约束由数学正确性或硬件限制决定。

1. 可除性： 将大小为 $N$ 的循环拆分为因子 $A$ 和 $B$ 时，乘积 $A \times B$ 通常必须等于 $N$（除非启用了循环填充）。
2. 资源限制： 在 GPU 上，每个块的最大线程数是固定的（例如，许多英伟达架构为 1024）。提议线程块大小为 2048 的调度是无效的，必须从空间中排除。
3. 内存容量： 需要共享内存分配超出硬件 L1 scratchpad 限制的平铺大小，将导致运行时错误。

软性约束

软性约束是用于引导搜索朝向高性能区域的启发式方法。虽然违反这些并不会使代码出错，但几乎肯定会导致性能不佳。

1. 向量长度： 在带有 AVX-512 的 CPU 上，使用不是向量寄存器大小倍数的长度（例如，16 个浮点数）进行向量化，通常不是最优的。
2. 展开限制： 过度循环展开会增加指令缓存的使用。编译器通常会在搜索空间定义中限制最大展开因子。

硬件绑定的参数化

在 TVM 或 Halide 等现代机器学习编译器中，搜索空间明确包含逻辑循环到物理硬件单元的映射。这与标准 CPU 循环优化有所不同。

对于 GPU 目标，搜索空间包含决定哪些循环级别映射到 blockIdx.x（CUDA 块索引）以及哪些映射到 threadIdx.x（CUDA 线程索引）的变量。

如果我们有一个平铺的循环嵌套：

1. 外层循环 $i$
2. 内层循环 $j$

搜索空间包含代表绑定策略的类别变量。

• 选项 A：将 $i$ 绑定到块，将 $j$ 绑定到线程。
• 选项 B：将 $i$ 绑定到线程，将 $j$ 绑定到本地寄存器（顺序执行）。

这一决定会根本性地改变内存访问模式。选项 A 利用流式多处理器之间的并行性，而选项 B 则可能用于线程局部累加。搜索空间定义必须列举这些结构可能性。

组合式增长

搜索空间的大小随算子复杂性呈指数级增长。这一现象使得除了简单的内核之外，穷举搜索都不可行。

考虑一个二维卷积层。编译器可能调优：

• 高度、宽度、输入通道、输出通道的平铺大小（4 个变量）。
• 每个平铺的拆分因子（例如，拆分成 2 或 3 层）。
• 循环顺序（4-6 个轴的排列）。
• 向量化长度。
• 展开因子。

即使每个变量的范围适中，总排列数也可达数十亿。

标准矩阵操作中，随着额外优化轴的引入，搜索空间扩展的估计。

模板与草图

为管理这种复杂性，现代方法通常采用“模板”或“草图”，而非原始的、非结构化搜索空间。

一个 模板 手动定义调度的骨架结构，但将具体整数值留作变量。例如，工程师可能编写一个调度模板，其中规定：“始终平铺输出通道并并行化外层循环，但让自动调优器决定平铺大小 $X$ 和展开因子 $Y$。”

这种混合方法将搜索空间限制在已知普遍高效的区域，仅依靠自动调优器进行数值参数的硬件特定调优。这能大幅减少搜索时间，同时仍能获得高性能。

在下一节，我们将了解成本模型如何分析这些候选，以预测性能，而无需在实际硬件上运行每一个。