论文精读（四）

Large Language Models for Compiler Optimization

• 组会报告的稿子，懒得整理了，可结合PPT查看。

• 今天我要讲的论文是“Large Language Models for Compiler Optimization”，是 meta 上个月放在 arxiv 上的一篇论文，我感觉比较有意思。

• 首先我们来看一下文章的背景

  • 大语言模型在学术界工业界都引起了极大的关注，下图是有关大语言模型的论文的增长趋势，可以说是爆炸式地增长。大语言模型在软件工程领域也是十分引人注目，包括代码生成，代码翻译，代码测试，大家平时如果有使用 chatgpt 的话应该有所体会。我之前旁听过汇丰银行来超算中心开的一次会，当时黄老师就问他们，“大语言模型在汇丰银行里有什么应用”，当时的回答也是这些。
  • 但是大模型在代码优化领域的应用还是一片空白，本文是第一个将大模型用在代码优化领域。
  • 我们再来看看基于机器学习的代码优化解决方案目前是什么情况，文章主要举了两个例子，一个是 MLGO，它使用了人为选取的特征来进行机器学习；另一个是 PrograML，它将程序转化为图再喂给图神经网络。不论是哪种方法，对于机器学习模型来说，输入都损失了信息，相比源代码是不完整的。从这个角度看，大语言模型是有优势的，因为它的输入可以是源代码本身。

• 接下来看一下论文的 Design，这张图是整个流程。大语言模型的输入 prompt 有自然语言和代码两部分，让大语言模型生成指令数最少的 pass 顺序，训练的时候会给出较优的 pass 顺序和优化后的代码，推理的时候会生成 pass 顺序，然后给编译器来实际执行优化。

  • 这篇论文的一个有意思的点就是把代码优化的落脚点放在了 pass list 来减小代码大小上，因为 pass 顺序不管怎么变，它的正确性是由编译器保证的，不会出现大量的错误，这样就看起来比较靠谱。
  • 这篇文章的另外一个有意思的点就是源代码使用了 LLVM-IR，而在过去大语言模型通常都使用 Python 作为输入。个人感觉设计上的一个比较大的弱点是它直接把 LLVM-IR 的指令条数当做代码大小，而不是编译得到的二进制程序的代码段大小。
  • 接下来我们来看一下实现上的一些细节。首先是 Prompt，不知道大家有没有了解过提示词工程，它有一个技巧，就是为了让大模型深度思考，我们要学会循循善诱。什么意思呢，大模型有点像人，给它命令时，过程和结果一样重要，当我们要求大模型得到一个困难的最终结果时，最好先让它思考一个中间问题或者高度相关的问题，这通常会提升它给出最终结果的质量。那本文就让大模型多思考了两个额外问题，一个是给出实施对应的 pass 前后的 LLVM-IR 指令数，另一个是优化后的 LLVM-IR。这些结果并不需要用在最终的部署环境中。此外，对于输入代码，还全洗了一遍来保证规范化，并且减少输入长度。
  • 然后我们来看一些模型的细节。这个模式的架构是一个 7B 的 LLaMa 2，是从头开始训练的，而不是使用预训练模型再微调，文章还给出了一个大致的训练数据输入的代码长度上限：2KB。
  • 然后是训练数据。训练数据的输入部分，即未优化的 LLVM-IR，主要从公开的代码库提取函数，并且针对性地生成一些测试代码。为什么训练数据主要是在函数级别，而不是单个源代码文件甚至整个项目？这是因为输入代码最大只能有 2KB。接着是训练数据的输出部分，即 pass 的顺序，文章实现了一个 autotuner 来自动地搜索最优的结果，搜索空间为 122 个 pass 和 6 个优化级别，经过实验发现 pass 列表通常包含 9 个 pass，所以搜索空间大约为$10^{18}$，平均每个程序编译了 37,424 遍，最后和-Oz相比指令数减少了$5.8%$。自动调优的方法比较暴力，是优化效果的天花板，但是非常耗时，本文希望利用大模型替换这种方法。

• 最后我们来看一下实验部分。

  • 首先是训练，本文用了64块V100，共训练620个GPU days，大概就是64张卡训练10天，这么多硬件资源实验还是有不小门槛。在验证集上，文章从三个角度比较了训练的结果，刚好是之前 Prompt 部分提到了要求大模型进行的几个工作。第一个是相比-Oz，生成的 pass 列表的指令数减小幅度，当然比不过自动调优，因为是理论上限，但是总体上效果还是不错的。第二个是指令数量的预测，它对输入代码指令条数的预测比较准确，但对生成代码条数的预测比较差。第三个是生成代码的预测，对比的对象为编译器生成的代码，最下面的是完全匹配的比例，绿色的是可以通过编译的比例，最上面的是 BLEU 的一个相似度衡量指标。大模型生成的代码在验证集上总体看起来还不错。
  • 我们接下来看一下大模型和当前一些 SOTA 实现的比较。文章找了不同领域的六个测试集，选取了3个baseline，一个是强化学习方法，一个是带 cost model 的迭代搜索，还一个是之前提到了自动调优器。这个图是和-Oz比较的结果，可以看到大模型泛化较好，在测试集上也是效果拔群。至于为什么这里另外两个 baseline 是负的，因为在预测的时候优化级别可能不是-Oz，所以文章进行了额外的实验，如果预测的优化级别不是-Oz，则把优化级别改为-Oz，进行额外的编译，选较好的结果，即使这样，可以看到大模型的结果还是好于另外两个。
  • 再来看一下一些结果的细节上，一个是大模型预测的各种 pass 出现的频率，右侧是预测的 pass 列表长度，大家可以看到各种熟悉的 pass。这是一个大模型优化效果比自动调优好的例子。
  • 接着是大模型给出的优化后代码的质量。之前提到大模型给出的优化后的代码相当一部分是根本不能用的，文章把各种编译错误进行了分类。大模型还是经常会说胡话。
  • 文章还进行了额外的实验，一个是改变模型训练数据的大小，观察到训练数据越多，效果越好；同时为了验证之前提到的 “Deep Understanging” 的能力，也训练了一个只要求生成 pass 序列，而不要生成额外内容的模型，相比下最终结果有所下降。
  • 文章对大模型进行单个 pass 的优化进行了评估，针对 60 个 pass， 每个 1w 个训练数据对。一共训练了 74 GPUdays。下图是结果，不同的 pass 间还是有不小差异。它还真的做了循环展开和循环折叠的实验，大概只有 60% 能过编译，但是最后面的循环融合正确性就很高。

• 最后文章对本文的工作进行了一些讨论，第一个是大模型输入的长度，本文只能处理2KB的程序，大家都知道大模型是$O^2$的，所以这个问题是所有大模型都存在的问题，可能需要业界共同来解决；第二个有关大模型的推理能力，这个是之前反复提到的，人类来把一个复杂的问题分解为多个问题，让大模型深度思考，这再次被证明是有益的；第三个是推理速度，目前比编译器慢两个数量级，但是比基于搜索的自动调优要块得多。

• 总结：工作很新，有一定参考价值，但是槽点也不少。