在物理硬件上评估数千个候选调度会在编译流程中造成明显瓶颈。虽然遗传算法或模拟退火求解器可以快速生成有效配置,但内核在GPU或TPU上的实际运行需要毫秒到秒级时间。当搜索空间包含数十亿种可能性时,硬件在环验证在计算上会变得难以负担。为弥合庞大搜索空间与有限硬件可用性之间的差异,自动调优器使用统计成本模型。成本模型作为硬件的高吞吐量替代品,无需执行即可预测特定调度的性能。这使搜索算法每秒能查询数千个候选的性能,从而筛选出最有潜力的配置进行实际硬件基准测试。替代模型的作用正式来说,我们将优化问题定义为从有效调度空间 $S$ 中找到一个调度 $s$,使其在硬件目标 $H$ 上的运行时间最小。真正的成本函数 $Cost(s, H)$ 计算起来很昂贵。成本模型提供了一个近似函数 $\hat{f}(s)$,满足:$$ \hat{f}(s) \approx Cost(s, H) $$搜索算法主要与 $\hat{f}(s)$ 进行交互。该工作流程采用主动学习方法。调优器最初不了解硬件对特定算子的响应。它随机选择一小批调度,在设备上进行测量,并使用得到的延迟数据来训练初始成本模型。随着调优的进行,模型会预测下一批候选的性能。最佳候选在硬件上进行测量,这些真实测量数据被添加到训练集中以更新模型。这个反馈循环不断提高模型在搜索空间中关键区域的准确性。digraph G { rankdir=TB; node [shape=box, style=filled, fillcolor="#e9ecef", fontname="Helvetica", fontsize=10]; edge [fontname="Helvetica", fontsize=9]; subgraph cluster_0 { label = "宿主机编译器"; style=dashed; color="#adb5bd"; SearchPolicy [label="搜索算法\n(遗传/模拟退火)", fillcolor="#d0bfff"]; CostModel [label="统计成本模型\n(XGBoost/LightGBM/MLP)", fillcolor="#a5d8ff"]; FeatureExtract [label="特征提取\n(AST -> 向量)", fillcolor="#bac8ff"]; } Hardware [label="硬件测量\n(GPU/加速器)", fillcolor="#ffc9c9"]; Dataset [label="训练数据集\n{(调度, 延迟)}", shape=cylinder, fillcolor="#ffe066"]; SearchPolicy -> FeatureExtract [label="生成候选"]; FeatureExtract -> CostModel [label="输入向量"]; CostModel -> SearchPolicy [label="预测分数"]; SearchPolicy -> Hardware [label="Top-k 候选"]; Hardware -> Dataset [label="运行时间"]; Dataset -> CostModel [label="再训练 (梯度更新)"]; }自动调优中的主动学习循环。成本模型作为过滤器,使搜索策略能够查看数百万个候选,同时只在物理硬件上运行最有潜力的那些。程序特征工程机器学习模型需要数值输入,但循环调度是一个结构化的语法对象。将抽象语法树 (AST) 或特定循环配置转换为特征向量,是构建有效成本模型的重要一步。统计特征提取在像 AutoTVM 这样的系统中,特征提取依靠专家定义的规则。编译器遍历循环嵌套的低级中间表示 (IR),并提取统计计数器。常见特征包括:内存访问模式: 连续与步进式内存访问的数量。这有助于模型预测缓存一致性和内存带宽利用率。算术强度: 浮点运算量与传输内存字节数之比。循环结构: 循环嵌套的深度、每个循环的范围以及展开因子。向量化: 向量指令的长度和 SIMD 通道的利用率。这些特征被平展成固定长度的向量,并输入到梯度提升决策树 (GBDT) 中,例如 XGBoost。树模型在此处特别有用,因为它们处理特征之间的非线性关系(例如,瓦片大小与缓存大小),而无需数据归一化。基于深度学习的提取Ansor (自动调度器) 等新型系统使用深度学习来自动学习特征表示。这些系统不是使用手工设计的计数器,而是将程序视为图或序列。例如,调度可以被编码为一棵树,其中每个节点代表一个循环或语句。Tree-LSTM (长短期记忆) 网络或图神经网络 (GNN) 处理这种结构以生成嵌入向量。这种嵌入捕捉到复杂的依赖关系,例如定义与其使用之间的距离,这在简单的统计计数器中可能会丢失。排序而非回归虽然目标是最小化延迟,但成本模型不一定需要精确预测以毫秒计的运行时间。更重要的是,模型能正确地对两个调度的相对性能进行排序。如果调度 A 在硬件上比调度 B 快,即使绝对值不准确,成本模型也应预测 A 的分数低于 B。因此,现代成本模型通常优化成对排序损失,而非标准均方误差 (MSE) 回归损失。给定两个调度 $s_i$ 和 $s_j$,其测量延迟分别为 $y_i$ 和 $y_j$,预测分数分别为 $\hat{y}_i$ 和 $\hat{y}_j$,RankNet 风格的损失函数将错误排序的概率降到最低。我们使用 sigmoid 函数 $\sigma$ 定义 $s_i$ 比 $s_j$ 快的概率为:$$ P_{ij} = \frac{1}{1 + e^{-(\hat{y}_i - \hat{y}_j)}} $$损失函数会惩罚预测概率与硬件上观察到的实际排序之间的差异。这种方法使模型对硬件测量中的噪声具有抵抗力,并将学习能力集中在区分好坏调度上,而不是学习 GPU 的精确物理特性。搜索空间效率使用成本模型带来的效率提升非常显著。卷积神经网络 (CNN) 层的典型调优会话可能涉及查看大小为 $10^{10}$ 的搜索空间。没有成本模型,随机搜索或简单启发式方法可能会停滞不前,只找到局部最优解。有了成本模型,搜索过程能够有效地为每次实际硬件测量“模拟”数千步。下表展示了不同搜索策略的收敛行为。{ "layout": { "title": "收敛速度:随机搜索与成本模型引导", "xaxis": { "title": "硬件测量试验次数", "showgrid": true, "gridcolor": "#dee2e6" }, "yaxis": { "title": "吞吐量 (GFLOPS)", "showgrid": true, "gridcolor": "#dee2e6" }, "plot_bgcolor": "white", "paper_bgcolor": "white", "width": 700, "height": 400, "legend": { "x": 0.7, "y": 0.1 } }, "data": [ { "x": [0, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500], "y": [100, 150, 180, 200, 210, 215, 220, 225, 230, 232, 235], "type": "scatter", "mode": "lines+markers", "name": "随机搜索", "line": {"color": "#adb5bd", "dash": "dash"} }, { "x": [0, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500], "y": [100, 400, 650, 800, 880, 920, 940, 950, 955, 960, 962], "type": "scatter", "mode": "lines+markers", "name": "XGBoost 成本模型", "line": {"color": "#339af0", "width": 3} } ] }调优收敛的比较。成本模型(蓝色)快速识别高性能区域,在 300 次试验内达到接近峰值的吞吐量,而随机搜索(灰色)改进缓慢。成本模型中的迁移学习标准成本模型的一个明显局限是它们通常是任务特定的。训练用于优化形状为 $224 \times 224$ 的 Conv2d 算子的模型,通常无法预测 MatMul 甚至不同形状 Conv2d 的性能。这要求自动调优器为神经网络中的每个新子图从头开始。为解决此问题,高级编译器采用迁移学习。通过在包含各种算子和调度的大量数据集上进行训练,一个基础模型能学习到代码效率的通用原则,例如非对齐内存访问的惩罚或寄存器分块的好处。在调优新算子时,系统会使用这个预训练模型开始,并用少量逻辑样本进行微调。这显著减少了“冷启动”时间,使调优器能够用更少的硬件测量找到最优调度。