深度学习算子优化的自动化始于对优化边界的规范说明。编译器理论将程序看作功能等价物多维范围内的某一点,而非静态产物。此范围即为搜索空间。它包含生成正确数学结果但执行策略不同的所有可能内核版本。构建此空间需要对前几章提及的转换原语进行参数化。我们不手动编写具有固定分块大小(例如 $32 \times 32$)的特定调度,而是将这些值作为变量公开。编译器随后将代码生成视为一个约束满足和最小化问题,其目标函数是执行时间。循环转换的参数化搜索空间的基础在于应用于中间表示(IR)的转换原语。对于标准张量操作,主要的变化维度是循环分块(分割)、排序和硬件绑定。以一个简单的密集矩阵乘法循环嵌套 $C[i, j] += A[i, k] \times B[k, j]$ 为例。为了定义搜索空间,我们将具体的整数常量替换为符号变量。分割因子应用循环分块时,范围为 $N$ 的循环会被分割成外部循环 $N_{outer}$ 和内部循环 $N_{inner}$。分割因子 $f$ 定义了内部循环的大小。对于 $f$ 的有效值集通常是 $N$ 的因子。但如果允许循环填充,则 $f$ 可以是任何小于或等于 $N$ 的整数。如果应用多级分块(例如,映射到 GPU 的线程块、warp 和线程级别),我们会引入多个分割因子。对于一个长度为 $L$ 并分割成 $k$ 级的循环,此维度的搜索空间成为整数分区集,满足:$$ \prod_{i=1}^{k} f_i = L $$循环排列循环分割后,它们嵌套的顺序决定了数据局部性。对于深度为 $d$ 的嵌套,有 $d!$ 种可能的排列。虽然一个简单的3循环嵌套只有6种排列,但三级分块的矩阵乘法会产生9个循环(每个 $i, j, k$ 轴一个外层、中层和内层循环)。这会得到 $9! = 362,880$ 种理论上的排序。我们立即利用依赖分析来剪枝此空间。只有保持数据依赖的排列才会被纳入有效搜索空间。组合与树形结构我们可以将搜索空间的构建想象成一棵决策树。每个节点表示一个转换选择(例如,“分割轴 $i$”),分支表示参数值。从根到叶的路径表示一个完全定义的候选调度。digraph G { rankdir=TB; node [fontname="Sans-Serif" shape=rect style=filled fillcolor="#e9ecef" color="#adb5bd"]; edge [color="#868e96"]; Root [label="初始循环嵌套 (i, j, k)" fillcolor="#a5d8ff"]; Split_I [label="分割轴 i" fillcolor="#bac8ff"]; Split_J [label="分割轴 j" fillcolor="#bac8ff"]; Factors_I1 [label="因子=[4, 8, 16...]" fillcolor="#eebefa"]; Factors_I2 [label="因子=[32, 64...]" fillcolor="#eebefa"]; Reorder1 [label="重排 (i_o, j_o, k, i_i, j_i)" fillcolor="#ffc9c9"]; Reorder2 [label="重排 (j_o, i_o, k, j_i, i_i)" fillcolor="#ffc9c9"]; Vectorize [label="向量化内部循环" fillcolor="#b2f2bb"]; Unroll [label="循环展开内部循环" fillcolor="#b2f2bb"]; Root -> Split_I; Root -> Split_J; Split_I -> Factors_I1; Split_I -> Factors_I2; Factors_I1 -> Reorder1; Factors_I2 -> Reorder2; Reorder1 -> Vectorize; Reorder2 -> Unroll; Vectorize -> Result1 [label="候选 A"]; Unroll -> Result2 [label="候选 B"]; Result1 [label="调度 A" shape=oval fillcolor="#fcc2d7"]; Result2 [label="调度 B" shape=oval fillcolor="#fcc2d7"]; }转换决策如何分支形成搜索空间的层次视图。每个层级都添加一个约束,缩小了可能性,直到形成一个具体的调度。硬件约束与有效性并非所有参数组合都能产生可执行代码。“有效”搜索空间是参数总笛卡尔积的一个子集,受硬件限制。资源限制: 在 GPU 上,映射到 threadIdx.x、threadIdx.y 和 threadIdx.z 的线程索引乘积不能超过每个块的最大线程数(通常为 1024)。一个分别以 64 和 32 为因子对维度 $i$ 和 $j$ 进行分块的调度将需要 $64 \times 32 = 2048$ 个线程,这使其无效。共享内存: 分块大小决定了提升到共享内存的内存块大小。如果计算出的内存占用超出每个流式多处理器(SM)可用的共享内存,内核将无法启动或占用率下降。向量化约束: 向量化循环要求内存访问模式连续,并且循环范围是向量长度的倍数。试图向量化非连续维度的调度会被丢弃。搜索机制必须有效筛选这些无效配置。在 Ansor (AutoTVM v2) 等框架中,这通常由草图生成规则来处理,这些规则只提出结构上有效的调度,然后通过随机抽样步骤验证资源约束。基于模板与基于生成的空间现代编译器中定义搜索空间边界有两种主要方法。基于模板(AutoTVM)在这种方法中,工程师编写一个通用调度“模板”,并将特定值标记为可调节参数。例如:$$ \text{cfg.define_split("tile_k", k, num_outputs=2)} $$搜索空间由用户的定义明确限定。编译器严格在提供的参数范围内优化。虽然这确保了优化结构是已知的(例如,我们肯定在使用分块和缓存),但它需要工程师的直觉,将最佳优化策略包含在模板中。如果用户忘记包含循环展开的参数,自动调优器将永远不会尝试它。基于生成(Ansor/MetaSchedule)新的系统会从计算图自动生成搜索空间。编译器分析算子的数学属性,并应用一组推导规则来构建空间。草图生成: 系统创建高层次的结构草图。一个草图可能表示“总是对归约轴进行分块”,而另一个则表示“内联此逐元素操作”。标注: 一旦选定草图,系统会随机标注具体细节,例如分块大小或展开因子。这种方法显著扩展了搜索空间。它消除了模板设计中的人为偏见,但增加了搜索的复杂程度。该算法必须遍历一个更大的范围,这使得我们在本章后面将研究的成本模型成为必需。组合爆炸搜索空间的大小随算子的复杂性和硬件层级的深度呈指数增长。对于针对 GPU 的 Conv2D 算子,我们必须决定:批次 ($N$)、高度 ($H$)、宽度 ($W$)、通道 ($C$)、核 ($K$)、滤波器 ($F$) 的分块因子。将这些分块映射到块 X/Y/Z 和线程 X/Y/Z。缓存读/写位置(全局到共享,共享到寄存器)。展开因子和向量化长度。即使积极剪枝无效配置,空间也容易包含数十亿个候选方案。{"layout": {"title": {"text": "搜索空间大小与优化深度对比(对数刻度)", "font": {"family": "Arial, sans-serif", "size": 18, "color": "#495057"}}, "xaxis": {"title": {"text": "分块层级数", "font": {"size": 14, "color": "#868e96"}}, "gridcolor": "#e9ecef"}, "yaxis": {"type": "log", "title": {"text": "有效调度(估计)", "font": {"size": 14, "color": "#868e96"}}, "gridcolor": "#e9ecef"}, "plot_bgcolor": "white", "paper_bgcolor": "white", "margin": {"t": 60, "l": 60, "r": 30, "b": 60}}, "data": [{"x": [1, 2, 3, 4, 5], "y": [120, 14400, 1728000, 207360000, 24883200000], "type": "bar", "marker": {"color": ["#a5d8ff", "#74c0fc", "#4dabf7", "#339af0", "#228be6"]}}]}随着我们增加内存层级(分块层级),候选调度的数量呈对数增长。此图表假设每个层级有适量的分割因子和重排选项。有效地在这个空间中找到最优解,需要的不仅仅是暴力搜索。我们定义此空间,使其包含最高性能的候选方案,同时保持其足够稀疏以便可以搜索。这种权衡是自动调优的核心挑战。通过将优化决策规范化为参数化范围,我们从手动启发式调优转向自动化、数据驱动的寻优。