学习用于通用代码优化的编译器世界模型，超越内置编译器优化标志和最先进方法！

https://arxiv.org/abs/2404.16077

有效的编译器代码优化在计算机和软件工程中至关重要。这些优化的成功主要取决于应用于代码的优化 pass 的选择和排序。

尽管大多数编译器依赖固定的 pass 序列，但目前为特定程序寻找最佳序列的方法要么采用不切实际的慢速搜索算法，要么使用难以泛化到训练期间未见过的代码的学习方法。

为了应对这些挑战，我们引入了 CompilerDream，这是首个用于通用代码优化的基于世界模型的方法。CompilerDream 具有带有奖励平滑技术的编译器世界模型，能够准确模拟优化过程。

基于此模型，可以通过值预测或直接优化序列生成来构建代码优化智能体。这些智能体在大规模程序数据集上进行训练，可作为跨不同应用场景和源代码语言的通用代码优化器。

我们的大量实验表明，CompilerDream 在自动调优方面具有强大的优化能力，在 CompilerGym 排行榜上领先。

更重要的是，经过大规模训练的编译器世界模型和智能体的零样本泛化能力在不同数据集上表现出色，在值预测和端到端代码优化两种场景下均超越了 LLVM 的内置优化和最先进的方法。

一、引言

代码优化在充分发挥软件和硬件潜力方面发挥着重要作用。开发人员期望有一种通用解决方案，能够将输入程序转换为语义等价但效率更高的版本，而无需人工干预。

编译器通过前端将源代码转换为中间表示（IR）、中间端优化器执行与语言和平台无关的 IR 优化，以及后端将 IR 转换为二进制代码来实现这一点（图 1）。

优化器通常以一系列 pass 的形式实现，对代码应用转换，其性能在很大程度上取决于这些优化 pass 的选择和顺序。标准编译器使用几组固定的优化序列来增强程序性能的特定方面，例如用于执行速度的-O1、-O2 和-O3，以及用于减小程序大小的-Os 和-Oz。

显然，鉴于程序和平台的巨大多样性，这些由编译器专家预先定义的现成策略在大多数情况下并非最优。因此，为特定程序自动优化 pass 序列可以比默认编译器设置带来显著的性能提升[26,65]。为了实用，这种算法必须在合理的时间内生成令人满意的 pass 序列，并处理各种各样的程序。然而，当前的研究往往无法同时满足这些要求。基于搜索的方法[8]实现了接近最优的结果，但每个程序需要数千次编译来验证优化结果，这使其不切实际。相比之下，机器学习方法通过直接预测优化序列或估计优化序列的结果来指导搜索，避免了这些耗时的编译器交互。

然而，这些基于学习的方法仍然存在牺牲一定最优性的风险，更重要的是，在跨可能超出训练样本的不同程序进行广泛泛化时面临显著瓶颈。

一系列机器学习应用[11,20,39,58]已经证明，在大规模数据集上训练高容量模型可以产生前所未有的性能。

先前关于编译器优化的研究也表明，使用具有不同程序的大规模数据集进行训练可能是有益的[17]，但该领域的普遍做法仍然是以每个程序的方式学习优化策略[34,61]，或者从仅包含几百个程序的相对较小的训练集和有限容量的模型中学习[37,52]，这阻碍了泛化能力。

本文关注 LLVM[41]的 pass 排序问题，这是编译器研究的一个长期挑战。

我们提出 CompilerDream，一种用于通用代码优化的基于世界模型的方法，能够处理各种各样的程序，这与大多数先前专注于狭窄领域特定方法的研究不同。

它学习一个准确的预测世界模型来模拟编译器执行，根据当前的 IR 状态和应用的 pass预测未来的 IR 状态和优化指标的改进。

我们认为，使用世界模型为编译器优化提供了以下优势：

首先，通过捕捉优化动态，世界模型将可泛化的知识融入我们的方法[4]，提高整体泛化性能。
其次，它取代了昂贵的编译器调用，显著减少了搜索和学习算法的计算开销，并促进了大规模训练。
此外，其高容量架构可以从广泛的数据集中提取更深入的见解。

具体来说，为了使世界模型更好地适应编译器优化，我们仔细研究了优化过程，并引入了奖励平滑技术来增强世界模型的训练并提高模拟准确性。基于这个准确的世界模型，CompilerDream 支持各种类型的优化智能体，包括值预测和强化学习。此外，我们精心策划了一个大规模的自然程序训练数据集，增强了世界模型和智能体对未见程序的泛化能力，使我们的方法能够预测更优的优化序列。

我们在一系列程序领域[17]和问题场景中展示了 CompilerDream 的有效性。

我们的测试领域包括涵盖基本算法的基准套件，以及生产级开源程序，如来自 C++ TensorFlow[1]和 OpenCV[13]库的目标文件。
即使不考虑泛化，CompilerDream 的强大优化能力也在 CompilerGym 自动调优排行榜上名列前茅。
对于通用代码优化，我们的大规模训练世界模型可以准确预测未见程序的 pass 序列结果。

配备世界模型的两种优化智能体在各自的场景中均优于内置的-Oz 标志和最先进的方法。

在值预测场景中，世界模型作为准确的值预测器，支持 superior action sequence selection。
在强化学习场景中，完全在世界模型中训练的智能体在单次试验中直接生成优化序列，在不同数据集上实现了更高效的代码。
这项工作的主要贡献如下：
我们提出了 CompilerDream，首个基于世界模型的代码优化方法。
我们引入了奖励平滑技术，促进了世界模型在编译器优化中的应用，有益于该领域的未来研究。
我们的方法支持多种优化智能体，包括值预测和强化学习智能体，统一了该领域的不同方法。
通过利用大规模 CodeContests 数据集和世界模型的泛化能力，我们在各种未见程序上实现了强大的性能。
在各种程序领域和问题场景（包括自动调优、值预测和端到端代码优化）的大量实验表明，CompilerDream 超越了最先进的方法。 :::

二、相关工作

编译中的一个关键挑战是确定应用哪些代码转换、如何应用（例如使用合适的参数）以及按什么顺序应用。

这涉及有效地搜索和评估众多选项，这一过程称为迭代编译[8]或自动调优[19]。然而，这种基于搜索的方法只能为一个特定程序找到良好的优化，无法泛化为编译器策略。这一局限性凸显了集成机器学习技术的重要性。

监督学习

开创性工作已深入探索监督机器学习，采用两种主要方法[43]。

第一种方法需要为每个训练程序进行广泛搜索，以确定最有效的优化序列，然后将其作为数据标签。早期示例[12]使用神经网络进行分支预测，一个更著名的工作是 MilepostGCC[25]，这是将机器学习集成到生产编译器 GCC 中的实际尝试。它采用在互联网上分布的大量程序数据集上训练的模型。
第二种方法旨在学习成本或性能函数，能够估计各种编译器选项的质量，从而无需编译和分析每个选项即可评估可能的选项[49,62]。Coreset-NVP[48]遵循这一方法，并实现了最先进的性能。

强化学习

最近的进展表明，强化学习（RL）技术在编译器优化中取得了进展，规避了收集最优标记数据的需要[40]。该技术已应用于优化单个编译启发式方法，例如内联[66]、循环变换[9,33]和图划分[55]。

与我们相关的几项工作，包括 AutoPhase[34]、CORL[52]和 POSET-RL[37]，已使用无模型 RL 探索完整的优化流水线，即 pass 排序问题，而未纳入世界模型。

世界模型

近似状态转移和奖励信号的世界模型[28,44]通常以两种方式使用：

它能够模拟“未见”的不可用或成本高昂的交互[30,38]；
作为辅助学习任务，它有助于捕获环境潜在结构的更好表示[4,50,53]。

据我们所知，我们是首个将世界模型引入代码优化的研究。在此背景下，世界模型可以充当编译器模拟器，近似 IR 转换，并消除对大量优化序列进行昂贵执行和分析的需要。

此外，通过与世界模型共享表示，策略可以更有效地泛化到未见程序。因此，可以实现基于模型的智能体[63]，与无模型方法相比，其提供了更高的样本效率和泛化能力。

三、方法

本节首先定义代码优化问题以及观测、动作和奖励的设计选择（第 3.1 节）。然后介绍 CompilerDream 世界模型的设计和训练技术（第 3.2 节）、两种优化智能体（第 3.3 节），以及大规模数据集策划的注意事项（第 3.4 节）。

3.1 pass 排序决策过程

如图 1 所示，编译器优化的一个关键问题是为给定程序找到最优的优化 pass 序列，这也被称为pass 排序问题。

在 POMDP 框架下，我们设计的观测和动作空间总结在表 1 中。

观测特征是为了效率和有效性而选择的。在初步实验中，我们发现 ProGraML[14]和 inst2vec[7]等复杂程序特征显著减慢了 CompilerDream 的速度，而性能提升有限。相比之下，专家设计的 Autophase 特征[34]利用领域知识，通过过滤无关细节来增强泛化能力。

我们通过将 56 维的 Autophase 特征向量与 42 维的动作直方图向量连接来构建观测，动作直方图记录了当前情节中每个动作被选择的次数。两个向量均进行了归一化：每个 Autophase 特征除以程序的初始总指令数，动作直方图按每情节动作限制（设置为 45）缩放。

在实验中，我们采用两种不同的动作空间来与基线方法保持一致。完整动作空间包含所有 124 个 LLVM 优化 pass，而从 Autophase[34]派生的精简动作空间包含 42 个动作。最初，该动作空间包含 45 个 LLVM pass，但由于最新 LLVM 版本的更新，CompilerGym 排除了 3 个 pass，剩下 42 个动作。这个精简动作空间被广泛使用，并在先前的研究[17,34]中被证明是有效的。

3.2 学习编译器世界模型

遵循先进的 Dreamer 方法[32]，我们构建了一个世界模型来学习编译器优化的 POMDP 过程，CompilerDream 的设计概览如图 2 所示。

编译器模拟

奖励平滑

我们发现，优化序列中的大多数优化 pass 不会改变程序 IR 指令计数，导致奖励稀疏。尽管如此，这些 pass 至关重要，因为它们可能修改关键属性，如指令顺序或替换，影响后续 pass 的有效性。此外，改进往往随时间饱和，因为大的非零奖励通常出现在序列早期，并随着优化进展而减少。

在图 3 中，我们展示了优化过程的一个典型情节：在 45 个动作中，只有 8 个产生非零奖励，其值随着优化进展而降低。这些特性为有效训练智能体带来了挑战，因为智能体从大多数动作中获得的指导很少，即使其中一些动作对最终结果至关重要。此外，世界模型可能难以准确预测奖励值，通常默认为零。这一挑战源于零奖励和非零奖励之间的严重类别不平衡，使得难以确定非零奖励的时机和幅度。

因此，我们应用奖励平滑来缓解情节内的奖励稀疏性和长尾分布。这通过向奖励添加指数衰减来实现：

3.3 学习代码优化智能体

基于世界模型，我们实现了两种智能体设计，将监督学习和强化学习方法统一到我们基于世界模型的方法中，用于编译器优化。

值预测智能体

这种智能体采用编译器优化中的经典方法，使用启发式方法指导搜索算法[48,49,62]。

它包括一个值预测模型来估计优化序列的效果，以及一个利用该模型识别最佳序列的搜索策略。

强化学习智能体

RL 是编译器优化的另一种流行方法，学习策略来有顺序地选择 pass。与依赖昂贵编译器交互的先前方法不同，CompilerDream 在世界模型模拟的编译器优化轨迹上训练 RL 智能体，显著提高了效率。

3.4 数据集策划

为了使 CompilerDream 方法能够有效泛化到未见情况（即零样本泛化），我们在准备训练数据集时确定了三个关键因素。

首先，数据集应反映自然性。在初步实验中，我们发现 Csmith[70]和 llvm-stress[41]等生成的代码数据集没有益处，甚至损害了对真实场景的泛化能力。
其次，数据集应足够大，以防止智能体过度拟合少量代码并无法泛化。
最后，从优化角度来看，代码必须具有高质量。我们需要具有复杂算法逻辑和优化潜力的数据，以帮助世界模型和智能体更好地理解问题的复杂性。像 AnghaBench[18]这样由从 GitHub 收集的数百万个人工编写的代码组成的数据集，通常过于简单，无法实现显著的优化改进。

因此，我们选择基于 AlphaCode[46]发布的 CodeContests 数据集构建训练数据集，该数据集包含超过 13,000 个编程竞赛问题，每个问题平均有数百个多种语言的解决方案。我们为每个训练问题最多抽取 10 个 C++解决方案，生成 110,240 个程序作为我们的训练数据，并为每个 100 个测试问题抽取一个解决方案作为验证数据。如表 2 所示，CodeContests 满足我们的所有标准，而其他数据集在一个或多个方面存在不足。

四、实验

我们在各种场景下进行了广泛的实验，将 CompilerDream 与最先进的方法进行比较。我们的研究旨在回答以下关键研究问题（RQs）：

RQ 1 最优性：CompilerDream 的编译器世界模型和优化智能体的联合训练过程能否发现更优的优化序列？
RQ 2 准确性：CompilerDream 的世界模型能否对真实优化过程进行准确模拟？
RQ 3 泛化性：通过世界模型学习的策略在各种未见程序上是否仍然有效？
RQ 4 有效性：CompilerDream 中采用的所有技术是否都能有效提升优化结果？

在以下章节中，我们首先展示 CompilerDream 作为强大的自动调优方法，通过发现更优的优化序列并利用准确的世界模型，在 CompilerGym 排行榜上领先（第 4.2 节），解决 RQ 1 和部分 RQ 2。然后，我们证明其值预测智能体（第 4.3 节）和 RL 智能体（第 4.4 节）在各种未见测试数据集上均优于最先进的方法，解决 RQ 2 和 RQ 3。最后，第 4.5 节通过消融研究回答 RQ 4。

4.1 评估

我们的实验侧重于代码尺寸减少，这对嵌入式系统等低资源硬件的应用有益。这一重点源于代码尺寸作为指标的实际优势：构建可编译的训练和测试数据集以及评估代码尺寸的优化性能成本较低。

指标

基准测试

我们主要在 CompilerGym 平台[17]的基准测试上评估我们的方法，包括 cBench[24]、CHStone[35]、MiBench[27]和 NASA 并行基准测试（NPB）[6]，以及来自开源项目的内核，如 BLAS[42]、Linux、OpenCV[13]和 TensorFlow[1]。不包括程序生成器[41,70]的合成基准测试，因为它们缺乏现实相关性。

我们遵循 CompilerGym 的标准数据划分。对于程序总数超过 100 的基准测试，我们使用前 50 个程序作为测试集，接下来的 50 个程序作为验证集，其余的作为训练集。这些训练和验证集仅用于域内训练。程序数量少于 100 的数据集不适用于域内训练，因此将其所有程序分配给测试集。划分后每个数据集的程序数量详见表 3。

此外，我们在一个名为 FormAI[64]的大规模数据集上评估 CompilerDream。FormAI 包含大量具有不同功能和编码风格的 AI 生成 C 程序。我们过滤掉无法编译成 CompilerGym 基准测试的代码，得到 109,016 个程序的测试集。受 FormAI 启发，我们还构建了一个由大型语言模型生成的 50 个 Objective-C 程序的数据集，以进一步评估 CompilerDream 对不同编程语言的泛化能力。为了进一步提升泛化能力，我们借鉴 FormAI [64] 的方法，使用 GPT-3.5 生成了一个包含 50 个独特 Objective-C 程序的数据集。我们使用与 FormAI 相同的提示，但添加了一条指令，要求 GPT 生成可以在 Clang 10.0.0 版本下直接编译且不使用 ARC（自动引用计数）的程序，以提高生成程序的编译通过率。使用 Clang 编译生成的程序，不包含任何第三方库，所有无法通过编译的程序均被丢弃。

4.2 自动调优：CompilerGym 排行榜

我们首先验证 CompilerDream 作为自动调优方法的能力，以展示其发现高质量优化序列的能力，支持我们实现卓越代码优化的目标。

实现

我们以 CompilerGym 排行榜任务[22]为目标，优化 23 个 cBench 程序的 pass 序列。遵循排行榜上其他方法的常见设置，我们使用 RL 智能体在除 ghostscript 之外的所有 cBench 程序上训练 CompilerDream，ghostscript 因体积庞大而显著减慢训练速度。动作空间设置为 LLVM 中所有 124 个动作的完整动作空间。CompilerDream 训练 25 小时，平均每个程序约 1 小时，与随机搜索算法的 3600 秒运行时间相当。除了评估单轮优化性能外，我们还测试了受排行榜领先方法启发的 CompilerDream+引导搜索。该方法利用 RL 智能体的策略指导随机搜索，将每个程序的搜索时间限制在 1 分钟内（详细信息见附录 B.6）。

结果

如表 4 所示，我们的 RL 智能体在世界模型模拟上训练，在单轮试验中在 cBench 上实现了平均 1.068× 的代码尺寸减少，尽管执行时间较长，但仍超越了随机搜索方法，即使考虑 CompilerDream 的训练时间也是如此。

这表明世界模型能够准确模拟优化过程并支持有效的智能体学习。此外，当与引导搜索结合时，CompilerDream 在使用更少运行时间的同时超越了排行榜的顶级方法。

4.3 通用值预测

在本节中，我们展示 CompilerDream 的大规模训练世界模型能够准确模拟未见程序的优化过程，从而构建出优于最先进方法的优秀值预测智能体。

实现

结果

表 5 展示了不同方法实现的几何平均减少。

尽管在与训练集不同的数据集上测试，CompilerDream 在三个数据集上超越了 Coreset-NVP，展示了卓越的编译器模拟和泛化能力。在 MiBench 上，CompilerDream 的表现接近 Oracle 基线（使用暴力搜索，作为问题的上限）。我们的方法和 Coreset-NVP 在 CHStone 上均实现了 1.101 的平均减少，接近 1.06 的上限。

4.4 基于 RL 的通用代码优化

我们最终展示 CompilerDream 的强化学习智能体能够在单次试验中端到端为训练期间未见的各种程序生成优化序列。这种场景反映了现实应用案例，其中优化算法只有有限的时间为任意程序计算最佳序列。因此，除非另有说明，本节中的所有方法在测试数据集上每个程序仅生成一个优化序列。

实现

我们使用 RL 智能体训练 CompilerDream，并将其与以下基线进行比较：随机 pass 序列、使用随机搜索的自动调优方法、LLVM 的-O0 和-Oz 标志，以及使用近端策略优化（PPO）[60]的最先进基于学习的方法[34]。随机搜索在与我们的 RL 智能体相当的时间预算内进行数百次试验。

结果

图 4 展示了 CompilerDream 的 RL 智能体实现的代码尺寸减少，通过 IR 指令计数减少的几何均值衡量。在没有域内训练的情况下，CompilerDream 除了两个基准测试外，在单次试验中均超越了-Oz，并且除了 BLAS 数据集外，始终优于 PPO。它在大多数数据集上也优于随机搜索，除了 Linux。BLAS 和 Linux 上各种方法的最小性能差异表明这些数据集已经高度优化。

此外，CompilerDream 的零样本泛化能力与域内训练相当或更好。这种优势在 NPB 数据集上尤为明显，其中数据稀疏性限制了域内智能体，而 CompilerDream 实现了额外 3%的代码尺寸减少。其强大的泛化能力扩展到新语言，如基于 Fortran 的 BLAS 和 NPB 数据集所示。在 AI 生成的 Objective-C 数据集上，CompilerDream 实现了平均 1.027× 的代码尺寸减少，在某些测试用例中高达 2.87×。

在大规模 FormAI 测试集（图 7）上，CompilerDream 超越了 PPO，在更多程序上匹配或超越-Oz，并实现了更高的优化水平，展示了其优越的一致性性能。

4.5 分析

与无模型方法的比较

我们进一步评估了基于模型的 CompilerDream（带 RL 智能体）与各种无模型方法（包括 PPO[60]、DQN[57]、A2C[56]、APEX[36]和 IMPALA[21,60]）的样本效率和零样本泛化能力。

图 6a 表明，虽然 PPO 是最强的无模型基线，但我们的 RL 智能体通过世界模型模拟训练，学习速度快 10 倍，且编译器交互更少。图 6b 进一步突出了其对未见基准测试的卓越泛化能力，支持我们的主张观点，即基于世界模型的智能体更好地捕捉编译器优化的动态，并增强对未见数据集的泛化能力。

与基于模型方法的比较

为了证明学习具有编译器深度表示的世界模型的必要性和有效性，我们在第 4.4 节描述的 RL 智能体设置下，将我们的方法与经典的基于模型的 RL 方法 MBPO[38]进行了比较。MBPO 为其动态模型采用标准 MLP，且不与演员或评论家共享潜在表示空间。

为了公平比较，我们使用与我们的方法相同的平滑因子对 MBPO 应用奖励平滑。结果如表 6 所示。尽管 MBPO 在某些任务上优于最佳无模型基线 PPO，但除了 MiBench（所有方法的性能差距较小）外，它始终不如 CompilerDream。

这些结果表明，【仅】基于模型的 RL 是不够的，更具表现力的世界模型和共享潜在空间对于更好的优化性能和对未见程序的泛化至关重要。

奖励平滑的效果

为了评估第 3.2 节所述奖励平滑技术的有效性，我们将使用完整 CompilerDream 方法训练的 RL 智能体的性能与未进行奖励平滑的变体进行了比较。

表 7 将这些结果与 PPO 进行了参考比较。结果表明，奖励平滑在所有数据集上持续提升优化性能，并使 CompilerDream 在 MiBench、Linux 和 OpenCV 上超越 PPO。

训练数据集的影响

为了评估 CodeContests 数据集对泛化能力的影响，我们将其与常用的 Csmith[70]数据集（一个由规则生成的大型 LLVM IR 数据集）进行了比较。

我们在两个数据集上使用 RL 智能体训练 CompilerDream，结果如图 6 所示，在所有五个测试数据集上，CodeContests 训练的 CompilerDream 显著优于 Csmith 训练的版本。这一优势在人工策划的 cBench[24]和 CHStone[35]数据集上尤为明显，表明 CodeContests 训练的模型能够更有效地泛化到人工编写的程序。

程序展示

图 5 展示了 cBench 中一个未见过的程序的预测优化轨迹，由我们学习的编译器世界模型预测。该模型成功预测了未来 IR 的数值特征，包括分支和块的计数，以及表示优化结果的未来奖励。

此示例体现了我们学习的编译器世界模型作为训练代码优化智能体的真实编译器环境可行替代方案的能力。

五、讨论

我们通过引入 CompilerDream 方法来解决基于学习的代码优化中的主要泛化挑战，该方法利用世界模型的模拟和泛化能力。

通过结合奖励平滑技术和大规模训练程序数据集，我们使世界模型能够有效训练以完成编译器优化任务。我们的方法支持两种优化智能体：值预测智能体和 RL 智能体。

实验结果表明，CompilerDream 在各种问题场景和程序数据集上均实现了卓越的优化性能，超越了内置的编译器优化标志和最先进的方法。

局限性

尽管我们的方法可以自然扩展到执行时间和目标文件尺寸减少等优化目标，但出于可扩展性和稳定性考虑，如第 4.1 节所述，我们仅专注于代码尺寸优化。

由于硬件和操作系统状态的影响，执行时间测量通常具有高方差和偏差，引入额外采样无法缓解的显著噪声。此外，现有开源框架（如 CompilerGym）缺乏对准确运行时测量的强大支持，使得难以获得有效的学习信号。

因此，本文【仅】专注于代码尺寸减少。

未来工作

有大量的探索空间，包括扩展训练数据集、扩大编译器世界模型规模、优化执行时间等多目标、以及利用深层专家知识或大型语言模型丰富特征和动作空间。此外，我们的方法可以支持更多类型的优化智能体，例如基于世界模型模拟的高级树 🌳 搜索算法。

参考文献

END

作者：THU
来源：NeuralTalk

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。