当前位置: 首页 > 科技观察

内存减少3%-7%!谷歌提出了机器学习框架MLGO

时间:2023-03-13 17:03:26 科技观察

用于编译器优化。随着现代计算机的诞生,出现了如何编译更快、更小的代码的问题。编译优化是成本效益比最高的优化方式,更好的代码优化可以显着降低大型数据中心应用的运行成本。编译代码的大小对于移动和嵌入式系统或部署在安全启动分区上的软件至关重要,因为编译的二进制文件必须符合严格的代码大小预算。随着该领域的进步,越来越复杂的启发式方法严重挤压了有限的系统空间,阻碍了维护和进一步改进。最近的研究表明,机器学习可以通过用机器学习策略代替复杂的启发式方法来释放编译器优化的更多机会。然而,在通用的工业级编译器中采用机器学习策略仍然是一个挑战。为了解决这个问题,谷歌两位资深工程师钱云迪和MirceaTrofin提出了“MLGO,一种机器学习引导的编译器优化框架”,这是第一个将机器学习技术系统集成到LLVM中的工业级通用框架,一种开源工业编译器基础架构,在构建关键任务的高性能软件中无处不在。论文地址:https://arxiv.org/pdf/2101.04808.pdfMLGO使用强化学习来训练神经网络做决策,取代了LLVM中的启发式。根据作者的描述,MLGO在LLVM上有两个优化:1)通过内联减少代码量;2)通过寄存器分配提高代码性能。这两种优化都在LLVM存储库中可用,并已部署到生产环境中。1MLGO是如何工作的?内联通过做出删除冗余代码的决定来帮助减少代码大小。在下面的示例中,调用函数foo()调用被调用函数bar(),后者又调用baz()。内联这两个调用站点将返回一个简单的foo()函数,这将减少代码大小。图注:内联通过去除冗余代码来减少代码大小在实际代码中,有成千上万的函数相互调用,从而形成一个调用图(Callgraph)。在内联阶段,编译器遍历所有caller-calleepairs的调用图,决定是否内联一个caller-calleepair。这是一个连续的决策过程,因为先前的内联决策会改变调用图,从而影响后续决策和最终结果。在上面的示例中,调用图foo()→bar()→baz()需要在两个边上都做出“是”的决定以保持代码大小下降。在MLGO之前,内联/非内联决策是通过启发式方法做出的,随着时间的推移,这种方法变得越来越难以改进。MLGO用机器学习模型代替启发式。在调用图的遍历过程中,编译器通过在图中输入相关特征(即输入)来寻求神经网络关于是否内联特定调用者-被调用者对的建议,并按顺序执行决策,直到遍历整个调用图.直到图形被调用。图例:内联期间的MLGO图,“#bbs”、“#users”和“callsiteheight”是调用者-被调用者对特征的实例MLGO使用策略梯度和进化策略算法训练在决策网络上执行RL。虽然没有关于最佳决策的基本事实,但在线RL使用经过训练的策略在训练和运行程序集之间进行迭代以收集数据并改进策略。特别是,编译器在内联阶段咨询模型以做出内联/非内联决定,给定当前正在训练的模型。编译后,它会生成顺序决策过程(状态、动作、奖励)的日志。然后将此日志传递给培训师以更新模型。重复此过程,直到获得满意的模型。图例:训练期间的编译器行为——编译器将源代码foo.cpp编译成目标文件foo.o,并执行一系列优化,其中之一是内联传递。经过训练的策略嵌入到编译器中,以在编译期间提供内联/非内联决策。与训练场景不同,此策略不生成日志。TensorFlow模型嵌入了XLAAOT,可将模型转换为可执行代码。这避免了TensorFlow运行时依赖性和开销,最大限度地减少了ML模型推理在编译时引入的额外时间和内存成本。图例:生产中的编译器行为我们在包含30k模块的大型内部包上训练了大型和小型内联策略。训练好的策略可以在编译其他软件时泛化,减少3%~7%的时间和内存开销。除了跨软件的通用性之外,跨时间的通用性也很重要,软件和编译器都在积极开发中,因此经过训练的策略需要在合理的时间内表现良好。三个月后,我们在同一组软件上评估了模型的性能,发现只有轻微的退化。图例:Inlinesizepolicy大小缩减百分比,x轴代表不同的软件,y轴代表缩减百分比。“Training”是训练模型的软件,“InfraX”是不同的内部软件包。MLGO的inlinesize-for-size训练已部署在Fuchsia上,Fuchsia是一种通用开源操作系统,旨在为以二进制大小为关键的各种硬件和软件生态系统提供支持。在这里,MLGO显示C++翻译单元的大小减少了6.3%。2寄存器分配作为一个通用框架,我们使用MLGO来改进寄存器分配(Registerallocation)流水线,从而提高代码在LLVM中的性能。寄存器分配解决了将物理寄存器分配给活动范围(即变量)的问题。随着代码的执行,不同的生命范围在不同的时间完成,为后续处理阶段释放寄存器。在下面的示例中,每个“加”和“乘”指令都要求所有操作数和结果都在物理寄存器中。有效范围x分配给绿色寄存器,并在蓝色或黄色寄存器的有效范围之前完成。x完成后,绿色寄存器变为可用并分配给活动范围t。在代码执行过程中,不同的活动范围在不同的时间完成,释放的寄存器用于后续处理阶段。在下面的示例中,每个“加”和“乘”指令都要求所有操作数和结果都在物理寄存器中。活动范围x分配给绿色寄存器,并在蓝色或黄色寄存器的有效范围之前完成。x完成后,绿色寄存器变为可用并分配给活动范围t。图例:寄存器分配示例当分配活动范围q时,没有可用的寄存器,因此寄存器分配流水线必须决定可以从其寄存器中“逐出”哪个活动范围以为q腾出空间。这被称为“字段驱逐”问题,是我们训练模型以替换原始启发式的决定。在这个例子中,它将z从黄色寄存器中逐出并将其分配给q和z的前半部分。我们现在考虑实际范围z的未分配下半部分。我们还有另一个冲突,这次活动范围t被驱逐和拆分,t的前半部分和z的后半部分最终使用绿色寄存器。Z的中间部分对应指令q=t*y,这里z没有用到,所以没有赋值给任何寄存器,它的值从黄色寄存器入栈,稍后再重新加载到绿色寄存器登记。t也会发生同样的情况。这会向代码添加额外的加载/存储指令,从而降低性能。寄存器分配算法的目标是尽可能减少这种低效率。这被用作指导RL策略训练的奖励。与内联大小策略类似,寄存器分配(regalloc-for-performance)策略是在谷歌内部的一个大包上训练的,可以跨不同的软件泛化,每秒在一组内部大型数据中心应用程序上查询的数量(QPS)增加了0.3%~1.5%。QPS的改进在部署后持续了几个月,显示了模型的通用性。3总结MLGO使用强化学习来训练神经网络进行决策,是一种替代复杂启发式的机器学习策略。作为一个通用的工业级框架,它将深入和广泛地进入更多的环境,而不仅仅是内联和寄存器分配。MLGO可以发展为:1)更深的,例如添加更多的特征,应用更好的RL算法;2)范围更广,适用于超越内联和重新分配的更多优化启发式。作者热衷于MLGO可以为编译器优化领域带来的可能性,并期待研究社区进一步采用它并做出未来的贡献。