大家好,我是 Chris Cummins,今天我要讲的是 CompilerGym。首先我想说,这是许多人的工作成果,我要感谢我在 Meta 的杰出合著者们的贡献。此外,这是一个开源项目,我也非常感谢社区做出的众多精彩贡献。
今天我们要讨论机器学习和编译器之间的关系。为什么我们关心这个?因为如果能让优化编译器生成更好的代码,每个人都能受益。你可以生产出更快的二进制文件,让用户更满意;还可以减少程序的能耗,从而节省大量电力和资金。更重要的是,如果你改进了编译器,那么该编译器生成的每一段代码都能得到提升。
当然,缺点也很明显:为了榨取尽可能多的性能而调整和优化编译器是一项非常困难的任务。这通常被认为是编译器专家领域的范畴,他们会花费数百甚至数千小时,通过试错和专业经验来调整启发式规则。但即便如此,也无法保证能获取所有潜在性能,往往还有性能被浪费掉。
因此,我们试图做的是将编译器优化的调整工作转化为数据科学和机器学习问题这样做的好处是,它不仅比人工调优做得更好,更重要的是我们将人类排除在循环之外。我们不再依赖编译器工程师为每次观察投入大量工时以获得最佳性能,从而提高了效率。
其工作原理是:你有一个编译器正在优化的程序(代码片段),需要将其表示为机器学习模型的输入。通常的做法是将该程序降低到中间表示(IR)层面,然后统计 IR 的属性,例如指令数量、循环嵌套深度或算术运算密度等。
想法是:如果你能收集足够多的不同程序,提取它们的特征,并确定每个程序中最佳的优化决策(即最佳参数),你就可以将这些作为训练数据。然后使用监督机器学习来学习一个模型,该模型可以将这些特征与产生最佳结果的优化决策关联起来。
如果成功做到这一点,你构建的模型就变成了你的优化启发式规则。将其嵌入编译器后,当新程序出现时,你只需将其映射到特征空间,模型就会预测它认为最佳的优化决策。这听起来很美好,但现实并不像我描述的那么简单。
事实上,如果你查看该领域的编译器论文,会发现它们通常在幕后隐藏了大量工作。那么编译器研究的实际情况是怎样的呢?假设我有一个利用机器学习改进特定编译器优化的绝妙想法,这很好,但我需要建立实验基础设施。
首先你需要一个编译器,这意味着你要弄清楚该特定编译器的 API。编译器极其复杂,通常有数百万行代码,这是一项令人望而生畏的任务。其次你需要数据集供学习使用,这意味着要收集一堆基准测试程序(benchmarks)。
这些程序通常打包松散且独立构建系统互不兼容,你经常不得不打补丁才能让它们正常运行。此外,你需要建立特征提取器的构建方法,这需要梳理既往工作以弄清其实现细节。这些代码往往不开源,或者即使开源也针对旧版编译器需要更新。
最后你需要将所有不同的脚本和插件粘合在一起,形成实验基础设施的一部分。当你完成这一切并运行实验时,会发现一堆 Bug 需要调试,整个过程充满繁琐的工作。希望在一个理想的世界中你能坚持下来尝试想法并获得结果,且结果是积极的。
这确实是一项繁重的工作。然而在其他机器学习领域(如自然语言处理或图像识别)中,通常有预打包的数据集、模型和框架工具库。这意味着你可以更快地通过快速迭代从想法得到结果,但这些在编译器研究社区中并不存在。
所以我们开始思考:如果我们能构建缺失的基础设施会怎样?这就是 CompilerGym 项目所做的。CompilerGym 有三个目标:第一,降低进入 AI 思维模式的门槛;第二,让研究人员尽可能轻松地开始探索真实的编译器问题。
我们真诚地相信,优化编译器的未来将通过机器学习优化的启发式规则来实现。第三我们的目标是提供一组通用基准测试,以便你可以在相同的编译器优化任务上比较不同的机器学习方法。
我们构建的东西看起来像这样:基本上,我们将编译器优化问题打包并通过强化学习循环暴露出来。这意味着它将一个问题(例如 LLVM 中的 pass 排序)封装在一个环境中,从而简化了研究流程。
环境是一个向用户暴露三个要素的小对象。第一个是观察(observation),即对代码状态的某种视图,例如表示程序的特征向量、图结构、源代码或程序的控制流图。第二个是奖励信号(reward signal),例如运行时时间或代码大小,用于指示特定优化决策的好坏。第三个是动作空间(action space),即该环境支持的优化选项集合,代表用户可做出的选择。
无论采用何种方法(不一定是强化学习),目标都是从动作空间中选择一个能带来最大累积奖励的动作。该过程不断迭代,直到环境告知已到达终止状态。CompilerGym 基于 OpenAI Gym 构建,这是一个用于强化学习的流行 Python 库。它非常灵活,允许用 Python(机器学习通用语言)编写智能体,同时也支持通过标准接口接入他人实现的库和智能体。
我们遵循“最小惊讶”原则,作为 Python 库,它返回标准的 Python 数据类型。处理数值特征向量时得到 NumPy 数组,处理图结构时得到 NetworkX 对象。接口是自描述的,用户无需了解特定编译器的内部实现即可交互,因此不需要相关领域的专业知识。
目前我们提供三个环境:首先是 LLVM 阶段排序(Phase Ordering),这是一个经典优化问题,旨在给定一组优化 passes(LLVM 有 124 个)的情况下,选择能最大化运行时或代码大小改进的 passes。第二个环境是命令行标志调优(Command Line Flag Tuning),从庞大的选项空间中挑选能生成最小程序的标志组合,动作空间为 500 维元组。
第三个环境是 Loop Tool Compiler,针对 CUDA 的循环嵌套代码生成任务。对于每个环境,我们暴露多种观察空间类型:可以请求程序的 IR 字符串、基于过往编译器研究的特征向量,或预训练的自然语言嵌入。我们还提供图表示以支持 GNN 等应用,并随时间不断扩充列表。
使用方式如下:首先实例化环境(如 Atari Breakout),获取初始观察,然后进行“步进”操作。每一步采取特定动作后,获得新观察、奖励以及是否结束的信号。示例中为简化演示随机选择动作(如挡板左右移动),实际应用中可接入 DQN 等深度强化学习方法。
CompilerGym 的设计使得只需更改环境名称(如改为 LLVM),代码即可操作编译器。此外,我们添加了处理编译任务的新 API:加载待优化程序以及导出优化后的程序,从而将编译器智能体集成到构建流程中。
由于实现了这一流行且简洁的接口,我们可以接入针对 Gym 优化的各种强化学习框架。例如使用 RLlib(Ray 库),它提供了多种最先进的算法实现。仅需约 15 行代码,即可使用近端策略优化(PPO)训练智能体进行 LLVM 阶段排序。RLlib 还提供 TensorBoard 日志记录和调试工具,开箱即用。
通过改变环境参数(如请求不同的观察空间),可以实验程序表示方式如何影响智能体的优化能力,或测试不同基准程序的泛化性。RLlib 内置了约 40 种强化学习算法,用户只需切换不同算法即可评估其性能差异。
智能体,或者这些不同算法——抱歉——在不同操作任务上的表现如何。现在并没有规定说你必须使用强化学习,Jim 并不要求你进行 RL(强化学习)。例如在这里,我们要展示的是如何将此与 Nevergrad 集成,这是一个黑盒优化器库,包含遗传算法以及大量用于无梯度优化的不同尖端技术。我们在这里做的是用几行代码来导航:这里是我们的函数,即 LLVM 阶段排序问题,去最小化它,去寻找针对该特定程序的最佳优化序列。或者你也可以从头开始编写自己的算法。
例如这里展示了一种贪心搜索的实现,而这里还有一个爬山搜索实现的示例。究竟是完全使用现成的强化学习工具或黑盒优化工具,还是深入其中并编写自己的算法,完全由你决定。CompilerGym 内置了许多额外的功能模块,我将快速概述其中几个重要的部分。首先是几乎 everything 都有命令行界面,这样你就可以在不编写任何代码的情况下进行大量此类优化问题的实验,这非常棒。
我们拥有的另一项功能是结果验证。一旦你运行了你的强化学习智能体,并且它表示在优化空间中进行了导航,环境随后可以通过自我验证来确保正确性。我的意思是检查程序的语义以验证它们是否保持一致,并验证程序没有发生变化,以确保你的编译器没有默默地破坏代码的语义。这对于发现编译器的错误或非确定性行为非常有用。
我们还拥有一组公共排行榜。一旦你构建了你的智能体或方法,并获得了结果集,将其提交到我们的 GitHub 仓库,我们将展示你的工作。你可以炫耀你的成果,并且我们会链接到你方法的详细文章。好的,我想在最后几分钟简要谈谈编译的架构,即内部机制。我不会深入细节,只是做一个高层概述:编译器的工作方式建立在分离关注点的架构之上,我们有前端和后端。
前端是一个大型 Python 库,也是用户与之交互的部分。它包含了所有面向用户的 API,例如管理数据集的方式、暴露实际 Gym 环境的方法、处理错误或从错误中恢复的支持,以及极化分布的任何支持都构建在这个前端内部。然后是命令行工具等组件。
后端则通过编译器服务接口与前端通信。这是一个相当小的接口,只有六个方法。但它以独立于编译器的形式暴露了为了展示编译器问题以供强化学习所需的操作。因此,你编写这些后端以针对特定编译器,例如我们拥有 LLVM 编译器服务,这是用 C++ 编写的编译二进制文件,针对 LLVM API 进行编译。它实际上是与编译器通信的部分。
其思想是你可以直接插入这些不同的后端,然后处理错误等样板代码以及所有面向用户的 API 都可以从前端重用。我们使其尽可能简单,以便你可以插入自己的编译器。我们在 GitHub 仓库中有一个示例,展示了如何添加新编译器到 CompilerGym。实际上它只占一页幻灯片的内容,我可以展示给你看:添加新编译器所需的最基本步骤。
你需要从名为 CompilationSession 的类继承,CompilationSession 抽象表示程序的单个编译实例。你需要声明动作空间列表,即编译器支持的决策选择;以及观察空间列表,即你的编译器知道如何计算的特征向量。然后你有三个方法:一个是 start(表示开始优化或编译一个新的程序,我们称之为基准测试);另一个操作是 apply_action(应用某个优化决策);下一个操作是 compute_observation(计算某些特征)。
如果你填好这五个空白,就可以获得一个编译器环境。这里展示的是 Python 实现或接口,也有等效的 C++ 接口,或者你可以使用你选择的语言。我们要往哪里去呢?就构建环境而言,我们的目标是稳步扩大搜索空间的大小。为此我们正在深入编译器内部。目前我们处理的是 somewhat modest 的问题,例如高层选择,如阶段排序(决定何时运行何种优化)。在这种情况下,我们的搜索空间仅为 10 的 200 次方。
接下来我们计划增加环境控制的粒度,开始关注单个循环级优化。一旦完成这一步,我们将进一步扩大范围到每一个指令级的优化决策。最终我们希望能够在编译器必须做出任何选择时,通过强化学习环境暴露出来。简而言之,我认为机器学习用于编译是一个有趣的领域,有很多产生影响力的工作空间,但存在较高的进入门槛。因此我们构建了 CompilerGym 来降低这一门槛。
这是一个库,它将入门门槛降低到了只需一个 pip install 命令。我们非常努力地使这个工具包做到“开箱即用”,也就是大家常说的功能齐全、无所不能。我们希望成为所有机器学习编译器研究的“一站式”商店。
我们目前正在积极开发它,我们的路线图非常详尽。请大家试用一下,并告诉我们您的体验如何。
非常感谢大家的时间。