编译器2.0：AI作为随机优化器

在HotChips上展示Jalapeño MLA内核及SemiAnalysis后续评论后，引发了大量讨论。作为OpenAI的硬件团队，我们只是勉强触及了这个“小金块”：AI正在编写我们的内核，而当它这样做时，我们并不真正需要逐行理解内核的功能。我们明显遗漏了：这种事情如何可能？与更传统的代码生成方法相比，正确的思考方式是什么？优化后的内核是否与未优化的同样可靠？

就我的背景而言，我在加速器编译器领域工作了十多年。我创立了XLA，这是一个优秀的编译器基础设施，拥有跨公司的杰出团队和努力。在过去两年多的时间里，在OpenAI，我一直在尝试重新构想编译器在AI时代应如何运作。新的编译器构想将利用现有优势，但不可否认，工具箱中增添了一个强大的新工具。

这将是一段旅程，但我希望阐明AI如何被用于自动化计算机程序改进，即优化编译。我确实认为，借助AI，我们可能会经历所谓的“编译器2.0”。AI在提议内容上受到的根本限制较少，其提议是模型训练和上下文的结果，这让我将其归类为随机优化器——这可能带来挑战，但正如我们将看到的，也是巨大优势的来源……

大量学术研究和行业应用已经朝这个方向发展，并迅速揭示AI在优化编译器领域的潜力，但我们现在正处于需要广泛概述的阶段。

背景

编译器接收程序，并输出这些程序的翻译版或改进版。

在输入和输出两端，程序都具有语义，这些语义告诉我们程序的含义、它们可能执行的操作，以及如何推理这些可能执行的操作。

我们这些从事编译器工作的人，往往将编译器视为纯函数——它们接收一个数据结构，并输出一个应具有相应语义的数据结构。

有时，我们的编译器侧重于“降级”或“翻译”。例如，它们可能接收C语言并输出x86-64汇编，这通常被认为是更低级别的。但往往，它们在此过程中所做的不仅仅是翻译……

在实践中，我们的编译器侧重于优化。它们可能接收一个表示程序的数据结构——用我们的术语来说，即“中间表示”（IR）——并尝试生成该程序的更优版本。有时，更优意味着运行所需周期更少，有时意味着减少不必要的代码，有时则意味着针对我们能够证明“必定为真”的程序特性进行特化（部分求值）。

现在，简要回顾一下，LLM最初是为了将人类文本从一种语言翻译成另一种语言而创建的。显然，翻译是它们的强项。通过日常任务中使用LLM，我们可以看到它们也能编写新的解决方案并改进现有方案。我们许多程序员也有过让LLM优化这段代码的经历，而它们确实能做到。（然而，我们需要确认它们是否正确优化了代码，这一点我们稍后会讨论！）这仅仅是为了强调，LLM具备我们在优化编译器中所寻求的能力。

优化与最优性

优化编译器，顾名思义，旨在通过某种目标（通常是执行时间）来提升所处理程序的最优性。在一般情况下，对于任意程序而言，这极其困难，以至于存在一个名为编译器工程师充分就业定理的定理。（我是在选择成为编译器工程师之后才发现这一点的，但它仍给我带来了安慰！）

“超级优化器”是优化编译器领域中一个令人惊叹的小分支。想象有一个给定的程序，我们可以通过语义来描述其功能。那么，具有相同语义的最优程序是什么？这正是超级优化器试图解决的问题，而它本质上是一个搜索问题……

假设我试图找到具有相同语义的最短程序，并且我有办法询问候选程序是否具有相同语义。理论上，我可以按目标顺序枚举每个程序，并选择具有相同语义的最小的那个。

然而，按目标顺序枚举每个程序听起来相当棘手。我最喜欢的学术论文之一，发表于2013年，题为“STOKE”（随机超级优化），提出了这样的疑问：“那么，如果我们只是反复随机调整程序，最终能否观察到最佳程序？”他们提出，通过随机游走（并借助我们最初级的机器学习朋友——马尔可夫链蒙特卡洛/梅特罗波利斯-黑斯廷斯算法），最终你会看到那个最优程序。

蒙特卡洛调整通常很笨拙（你随机选择一个调整），但速度很快。而LLM非常聪明（包含大量推理标记），但相对较慢。

如果我们不是用那种笨拙/快速的蒙特卡洛调整，而是让LLM来找出程序优化的方向呢？我们将拥有一个极其智能的随机优化器，引导我们的程序穿越优化过的程序空间。

优化的直觉

让我们退一步思考。想想你认识的最能体现“对代码片段进行极致优化”的人。简而言之，我们称他们为“优化狂人奥利”。奥利可能对哪些代码调整能产生效果有直觉。奥利可能会尝试一些方法看是否有效，如果不行，就回滚再试别的。但他们对于可能实现哪些优化，以及如何可能超越编译器，有着自己的直觉。

奥利拥有的这些直觉往往超出了编译器的能力范围。尽管现代优化编译器在结果上令人印象深刻，但它们基于相当简单的规则和启发式方法。用专业术语来说，它们基于局部数据流变换的思想，运行至固定点。我们还对考虑的因素进行阶段排序；也就是说，我们构建编译器流水线，先考虑A再考虑B，而不是综合的AB问题。调度器和寄存器分配器是这一问题的著名例子，许多博士研究尝试解决调度器与寄存器分配器结合的问题（以获取消除阶段排序的好处），但在实践中却难以实现。

这正是奥利专业知识的价值所在。奥利常常知道如何用针对特定情况的启发式方法平衡几个NP完全问题。因此，有更多针对具体情境的上下文意识和敏感性。奥利还能运用优化编译器可能无法有效应用的技术，尤其是组合使用时，比如提取代码、定制ABI或变换以启用向量化，以及其他让我们抱怨“真希望编译器有办法直接做到这个……”的各种技巧。

现在考虑一下，AI凭借其拥有的任何推理能力，或许能充当一个迷你版的Ollie。它可能没有那种对最终成果的直觉匹配，但它对什么可能盈利有模糊的感知，并且它可以进行很多，很多次尝试射门。

采用这种方法，与STOKE论文不同，我们无法保证随着时间趋向无穷，能看到最优程序，但因为AI拥有更接近人类的推理能力，它实际上能在单位时间内获得显著的人类式进展。

回归主题：MLA内核

让我先说明：我不知道AI为我们的Jalapeño MLA内核生成了什么样的低级代码，但我确实知道如何输入MLA的numpy代码。

在我之前工作的XLA编译器中，我们会将这些numpy操作融合成块，然后使用一个名为发射器的元程序将其降级为循环、指令和更底层的原语。

当XLA编译器/发射器程序这样做时，我不需要关心后端生成了什么汇编代码。对于我们的随机优化器，AI在概念上取代了发射器元程序的位置——它既进行降级又进行优化，我们可以要求它进一步优化，直至接近理论峰值性能。

我希望这能清楚地说明AI插在哪里，以及它如何与现有优化编译器系统中的组件相类比。思考这一点也很有帮助：我们现在认为的“汇编代码”层级正在上移。当你输入普通的C++并以-O3（典型的最高优化级别）编译时，即使你理解输入的C++，你也不会期望理解输出的汇编。我们在这里做的是类似的事情，但输入规格是更高级、更数学化的。

现在，一个关键问题是如何验证我们从AI那里得到的程序确实与更高级别的描述/numpy等价。这一验证机制确立了AI随机优化过程的可靠性。我预计未来的博客文章可能会对此进行更详细的探讨，但就目前而言，足以说明语义等价性测试是可行的，并且我们确实在这样做。加速器程序特别适合强健、完整的契约，我们可以验证这些契约“正是AI优化程序所执行的内容”，因为它们在广泛背景下相当数学化且以数据流为导向。

请注意，该领域的许多相关技术是由程序合成子领域的努力开创的。优化编译器会说：“这里有一个具有语义的程序，让它变得更好但保持语义等价！”而程序合成则说：“我相信存在一个具有这些语义的程序，请尝试找到你能找到的最佳版本。”程序合成比优化编译更难，但受根本性约束也更少。这实际上就是像Ollie这样的人在超越优化编译器时所做的事情，也是AI现在可以帮助我们自动化的内容。AI可以从原始程序中汲取“灵感”，但不必仅在其上进行微小的局部变换。经典优化编译器不会看到“哦，你写了一个冒泡排序”，然后通过理解契约将其切换为快速排序，但Ollie和AI都能做到这一点。这就是为什么我们更处于程序合成与随机优化的范畴，而非经典优化编译器的范畴。

这一切汇聚于一个事实：你可以从一个“与numpy相差不远”的起点开始，等待48小时，然后得到一个语义相同的优化内核，正如我们在HotChips演讲中所展示的：

如幻灯片所述，在我们的机器上，我们常常能观察到AI在性能上超越人类专家——即便是在我们认为调优得相当不错的kernel上。往往仍有相当可观的可实现百分比提升空间，这仅仅是因为可能需要探索的各种组合/排列方式繁多。对于人类性能工程师而言，这些探索往往繁琐到难以处理。

回顾与结论

归根结底，编译器只是一个函数。我们将程序输入该函数，得到的是程序的一个更优版本。我们期望输出的程序与输入的程序具有相同的语义。

传统优化编译器通过数据流规则和启发式方法使程序更优。这些方法的来源完全可理解，但可能受限于它们能执行的变换。

相比之下，AI作为一种随机优化器，只需深思熟虑并输出结果。其变换不受根本限制，使其更类似于我们的人类专家优化器。我们确实需要方法来验证AI输出的程序是否健全，并实现我们输入的相同语义，而我们已经具备了这些方法。这种AI优化特别适用于具有强契约的数学运算。这些契约避免了逐行理解kernel的需求。

这就是我们如何得到AI生成的MLA kernel的！