如何编译更快、更小的代码的问题随着现代计算机的诞生而出现。更好的代码优化可以显著降低大型数据中心应用程序的运营成本。编译代码的大小对于移动和嵌入式系统或部署在安全启动分区上的软件最为重要，因为编译后的二进制文件必须符合严格的代码大小预算。随着该领域的进步，越来越复杂的启发式算法严重压缩了空间，阻碍了维护和进一步改进。
最近的研究表明，机器学习 (ML) 可以通过用 ML 策略取代复杂的启发式方法，为编译器优化带来更多机会。然而，在通用的、行业强度的编译器中采用 ML 仍然是一个挑战。
为了解决这个问题，我们推出了“ MLGO：机器学习引导编译器优化框架”，这是第一个在LLVM（一种用于构建任务关键型高性能软件的开源工业编译器基础设施）中系统地集成 ML 技术的工业级通用框架。MLGO 使用强化学习 (RL) 来训练神经网络，以做出可以取代 LLVM 中启发式方法的决策。我们描述了两种针对 LLVM 的 MLGO 优化：1) 使用内联减少代码大小；2) 使用寄存器分配 (regalloc) 提高代码性能。这两种优化都可以在LLVM 存储库中找到，并且已经部署到生产中。
MLGO 如何工作？以内联大小为例
内联有助于减少代码大小，因为它可以做出删除冗余代码的决定。在下面的示例中，调用函数调用foo()被调用函数bar()，被调用函数本身又调用。内联两个调用点会返回一个可减少代码大小的baz()简单函数。foo()
内联通过删除冗余代码来减少代码大小。
在实际代码中，有数千个函数相互调用，因此构成了一个调用图。在内联阶段，编译器遍历所有调用者-被调用者对的调用图，并决定是否内联调用者-被调用者对。这是一个顺序决策过程，因为先前的内联决策将改变调用图，影响后续决策和最终结果。在上面的例子中，调用图foo()→ bar()→baz()需要在两端都做出“是”决策才能减少代码大小。
在 MLGO 出现之前，内联/非内联决策是通过启发式方法做出的，随着时间的推移，这种方法变得越来越难以改进。MLGO 用 ML 模型替代了启发式方法。在调用图遍历期间，编译器通过从图中输入相关特征（即输入）来向神经网络寻求是否内联特定调用者-被调用者对的建议，并按顺序执行决策，直到遍历整个调用图。
内联期间 MLGO 的说明。“#bbs”、“#users”和“callsite height”是调用者-被调用者对特征的示例。
MLGO 使用策略梯度和进化策略算法 ，通过强化学习训练决策网络（策略）。虽然没有关于最佳决策的基本事实，但在线强化学习会在训练和使用训练策略运行编译之间进行迭代，以收集数据并改进策略。具体来说，给定当前正在训练的模型，编译器会在内联阶段参考模型进行内联/非内联决策。编译完成后，它会生成顺序决策过程（状态、动作、奖励）的日志。然后将日志传递给训练器以更新模型。这个过程重复进行，直到我们获得满意的模型。
训练期间的编译器行为。编译器使用一系列优化过程将源代码编译foo.cpp为目标文件，其中之一是内联过程。foo.o
然后将训练好的策略嵌入到编译器中，以在编译期间提供内联/非内联决策。与训练场景不同，该策略不会产生日志。TensorFlow模型嵌入了XLA AOT ，它将模型转换为可执行代码。这避免了TensorFlow运行时依赖性和开销，最大限度地减少了编译时 ML 模型推理引入的额外时间和内存成本。
生产中的编译器行为。
我们在一个包含 30k 个模块的大型内部软件包上训练了内联大小策略。训练后的策略在应用于编译其他软件时具有可推广性，并且实现了 3% 至 7% 的大小缩减。除了跨软件的可推广性之外，跨时间的可推广性也很重要 — 软件和编译器都在积极开发中，因此训练后的策略需要在合理的时间内保持良好的性能。三个月后，我们在同一组软件上评估了模型的性能，发现性能仅略有下降。
内联策略大小减少百分比。x 轴表示不同的软件，y 轴表示大小减少百分比。“Training”是用于训练模型的软件，“Infra[1|2|3]”是不同的内部软件包。
MLGO 内联大小训练已部署在Fuchsia上，Fuchsia 是一种通用开源操作系统，旨在为多样化的硬件和软件生态系统提供支持，而二进制大小对于这些生态系统至关重要。在这里，MLGO 显示 C++ 翻译单元的大小减少了 6.3%。
寄存器分配（为了性能）
作为一个通用框架，我们使用 MLGO 来改进寄存器分配过程，从而提高 LLVM 中的代码性能。寄存器分配解决了将物理寄存器分配给有效范围（即变量）的问题。
在代码执行过程中，不同的有效范围在不同时间完成，从而释放寄存器以供后续处理阶段使用。在下面的示例中，每个“加”和“乘”指令都要求所有操作数和结果都位于物理寄存器中。有效范围x被分配给绿色寄存器，并在蓝色或黄色寄存器中的有效范围之前完成。x 完成后，绿色寄存器可用并被分配给有效范围t。
寄存器分配示例。
当需要分配有效范围q时，没有可用的寄存器，因此寄存器分配过程必须决定哪个（如果有的话）有效范围可以从其寄存器中“逐出”，以便为q腾出空间。这被称为“有效范围逐出”问题，我们针对这一问题训练模型以取代原始启发式方法。在这个特定示例中，它从黄色寄存器中逐出z ，并将其分配给q和z的前半部分。
我们现在考虑未分配的有效范围z 的 后半部分。我们再次遇到冲突，这次有效范围t被逐出并分割，t的前半部分和z的最后一部分最终使用绿色寄存器。z 的中间部分对应于指令q = t * y，其中z未被使用，因此未分配给任何寄存器，其值从黄色寄存器存储在堆栈中，稍后重新加载到绿色寄存器。t 也是如此。这会给代码添加额外的加载/存储指令并降低性能。寄存器分配算法的目标是尽可能减少这种低效率。这被用作指导 RL 策略训练的奖励。
与内联大小策略类似，寄存器分配 (regalloc-for-performance) 策略是在 Google 的大型内部软件包上进行训练的，并且可以推广到不同的软件中，在一组内部大型数据中心应用程序上，每秒查询数(QPS) 提高了 0.3% 至 1.5%。部署后，QPS 改进持续了数月，表明该模型在时间范围内具有普遍性。
结论和未来工作
我们提出了 MLGO，这是一个将 ML 技术系统地集成到工业编译器 LLVM 中的框架。MLGO 是一个通用框架，可以扩展为：1) 更深，例如，添加更多功能，并应用更好的 RL 算法；2) 更广泛，通过将其应用于内联和 regalloc 之外的更多优化启发式方法。我们对 MLGO 为编译器优化领域带来的可能性充满热情，并期待它进一步被采用，并期待研究界未来做出贡献。
亲自尝试
查看github 上的开源端到端数据收集和训练解决方案以及使用策略梯度训练内联尺寸策略的 演示。
致谢
我们要感谢 MLGO 的贡献者和合作者 Eugene Brevdo、Jacob Hegna、Gaurav Jain、David Li、Zinan Lin、Kshiteej Mahajan、Jack Morris、Girish Mururu、Jin Xin Ng、Robert Ormandi、Easwaran Raman、Ondrej Sykora、Maruf Zaber 和 Weiye Zhao。我们还要感谢 Petr Hosek、Yuqian Li、Roland McGrath 和 Haowei Wu 对我们的信任，并作为 MLGO 的第一位客户在 Fuchsia 中部署 MLGO；感谢 David Blaikie、Eric Christopher、Brooks Moses 和 Jordan Rupprecht 帮助在 Google 内部大型数据中心应用程序中部署 MLGO；并感谢 Ed Chi 和 Tipp Moseley 的领导支持。
评论