随着现代计算机的诞生，如何编译更快、更小代码的问题也随之出现。更好的代码优化可以显着降低大型数据中心应用程序的运营成本。编译代码的大小对于部署在安全引导分区上的移动和嵌入式系统或软件最为重要，其中编译的二进制文件必须符合严格的代码大小预算。随着该领域的进步，空间已经被越来越复杂的启发式方法严重挤压，阻碍了维护和进一步的改进。
最近的研究表明，机器学习 (ML) 可以通过用 ML 策略替换复杂的启发式算法来解锁更多编译器优化的机会。然而，在通用、行业实力的编译器中采用 ML 仍然是一个挑战。
为了解决这个问题，我们介绍了“ MLGO：机器学习引导编译器优化框架”，这是第一个在LLVM（一种开源工业编译器基础设施，普遍用于构建关键任务、高-性能软件）。MLGO 使用强化学习 (RL) 来训练神经网络以做出可以替代 LLVM 中的启发式的决策。我们描述了 LLVM 的两个 MLGO 优化：
- 通过内联减少代码大小；、
- 通过寄存器分配（regalloc）提高代码性能。
这两种优化都在LLVM 存储库中可用，并且已经部署在生产环境中。
MLGO 是如何工作的？
内联通过做出能够删除冗余代码的决策来帮助减少代码大小。在下面的示例中，调用者函数foo()调用被调用者函数bar()，而被调用者函数本身又调用baz(). 内联两个调用点会返回一个简单的foo()函数来减少代码大小。
在实际代码中，有数千个函数相互调用，因此构成了一个调用图。在内联阶段，编译器遍历所有调用者-被调用者对上的调用图，并决定是否内联调用者-被调用者对。这是一个顺序决策过程，因为之前的内联决策会改变调用图，影响以后的决策和最终结果。在上面的示例中，调用图foo()→ bar()→baz()需要在两条边上做出“是”决定，以减少代码大小。
在 MLGO 之前，内联/非内联决定是由启发式做出的，随着时间的推移，它变得越来越难以改进。MLGO 用 ML 模型代替了启发式算法。在调用图遍历期间，编译器通过从图中输入相关特征（即输入）来向神经网络寻求是否内联特定调用者-被调用者对的建议，并按顺序执行决策，直到遍历整个调用图.
MLGO 使用策略梯度和进化策略算法通过 RL 训练决策网络（策略）。虽然没有关于最佳决策的基本事实，但在线 RL 在训练和运行编译之间迭代，使用经过训练的策略来收集数据并改进策略。特别是，给定当前正在训练的模型，编译器在内联阶段参考模型进行内联/非内联决策。编译完成后，它会生成顺序决策过程（状态、动作、奖励）的日志。然后将日志传递给训练器以更新模型。重复这个过程，直到我们得到一个满意的模型。
然后将经过训练的策略嵌入到编译器中，以在编译期间提供内联/非内联决策。与训练场景不同，该策略不会生成日志。TensorFlow模型嵌入了XLA AOT ，可将模型转换为可执行代码。这避免了 TensorFlow 运行时依赖性和开销，最大限度地减少了 ML 模型推理在编译时引入的额外时间和内存成本。
我们在包含 30k 个模块的大型内部软件包上训练了 inlining-for-size 策略。训练后的策略在应用于编译其他软件时是可推广的，并实现了 3% ~ 7% 的大小减少。除了跨软件的通用性之外，跨时间的通用性也很重要——软件和编译器都在积极开发中，因此经过训练的策略需要在合理的时间内保持良好的性能。三个月后，我们在同一组软件上评估了模型的性能，发现只有轻微的退化。
MLGO inlining-for-size 训练已部署在Fuchsia上——这是一种通用开源操作系统，旨在为多样化的硬件和软件生态系统提供动力，其中二进制大小至关重要。在这里，MLGO 显示 C++ 翻译单元的大小减少了 6.3%。
Register-Allocation (for performance)
作为一个通用框架，我们使用 MLGO 来改进寄存器分配 pass，从而提高 LLVM 中的代码性能。寄存器分配解决了将物理寄存器分配给有效范围（即变量）的问题。
随着代码的执行，不同的生命周期在不同的时间完成，释放寄存器供后续处理阶段使用。在下面的示例中，每个“加法”和“乘法”指令都要求所有操作数和结果都在物理寄存器中。生存范围x分配给绿色寄存器，并在蓝色或黄色寄存器中的生存范围之前完成。x完成后，绿色寄存器变为可用并分配给生存范围t。
当需要分配生存范围q时，没有可用的寄存器，因此寄存器分配过程必须决定可以从其寄存器中“驱逐”哪个（如果有）生存范围以为q腾出空间。这被称为“生存范围驱逐”问题，并且是我们训练模型以替换原始启发式的决定。在这个特定示例中，它从黄色寄存器中驱逐z，并将其分配给q和z的前半部分。
我们现在考虑未分配的生存范围z 的后半部分。我们又发生了冲突，这一次生存范围t被驱逐和分割，并且t的前半部分和z的最后部分最终使用绿色寄存器。z的中间部分对应于指令q = t * y，其中z没有被使用，所以它没有分配给任何寄存器，它的值从黄色寄存器存储在堆栈中，后来被重新加载到绿色登记。同样的事情发生在t. 这会为代码添加额外的加载/存储指令并降低性能。寄存器分配算法的目标是尽可能减少这种低效率。这被用作指导 RL 策略训练的奖励。
与 inlining-for-size 策略类似，寄存器分配（regalloc-for-performance）策略是在一个大型 Google 内部软件包上进行训练的，并且可以在不同的软件中推广，每秒查询数提高 0.3% ~1.5% （ QPS）在一组内部大型数据中心应用程序上。QPS 改进在其部署后持续了数月，显示了该模型在整个时间范围内的普遍性。
结论和未来启发
我们提出了 MLGO，这是一个用于在工业编译器 LLVM 中系统地集成 ML 技术的框架。MLGO 是一个通用框架，可以扩展为以下几点：
- 更深，例如，添加更多特征，并应用更好的 RL 算法；
- 更广泛，通过将其应用于内联和 regalloc 之外的更多优化启发式算法。
我们对 MLGO 可以为编译器优化领域带来的可能性充满热情，并期待它的进一步采用和研究界的未来贡献。
评论 (0)