结合MLIR和PyTorch编译生态
允中 发自 凹非寺
量子位 | 公众号 QbitAI
Buddy Compiler 端到端 LLaMA2-7B 推理示例已经合并到 buddy-mlir仓库[1]主线。
我们在 Buddy Compiler 的前端部分实现了面向 TorchDynamo 的第三方编译器，从而结合了 MLIR 和 PyTorch 的编译生态。
目前，前端部分可以覆盖 LLaMA 计算图，转换到 MLIR 后我们集成了部分向量化和并行优化，并在 AVX512 平台上进行了测试。整个推理过程可以跑通但还需要大量优化。以下是相关链接和现状：
- [E2E] Buddy Compiler 端到端 LLaMA2-7B 推理示例[2]
- [E2E] 上述端到端推理示例目的是展示编译栈设计，并非完备的 LLaMA 问答工具
- [Frontend] Buddy Dynamo Compiler[3]
- [Midend] 集成面向矩阵乘法的向量化以及面向循环的并行优化
- [Backend] 端到端示例在 X86 AVX512 机器上进行测试（Ubuntu 22.04）
- [WIP] 开发并集成各种优化（现在速度太慢）
- [WIP] 在多种 CPU 后端进行测试（Arm Neon, RVV, etc.）
- [WIP] 对接多种硬件平台（GPU, Gemmini Accelerator, etc.）
- [WIP] 增加前端覆盖程度（Stable Diffusion, CLIP, Whisper, etc.）
概述
AI 大模型的爆发为软硬件设计带来了新的抓手和机会。随着模型的规模、类型、模态的增加和发散，单一软硬件技术栈覆盖各种场景的能力越来越受限。这种趋势也加深了我对软硬件生态重要性的认识。在我看来，整个生态的最重要的三条设计原则如下（重要程度由高到低）：
- 技术路线标准化（生态的立足之本）
- 上手门槛低（足够多的贡献者和玩家是生态繁荣的关键）
- 优化上限高（决定了生态的发展潜力）
我们基于以上思考打造了 Buddy Compiler， 致力于实现软硬件协同设计生态。我们的目标是实现从领域特定编程语言（DSL）到领域特定硬件架构（DSA）的编译流程和协同设计。
本文将介绍如何使用 Buddy Compiler 完成 LLaMA 2 的端到端推理。同时，我们也会分享 Buddy Compiler 的整体设计理念和未来的规划。具体的构建流程请参阅此处的文档[4]（大模型的构建需要十足的耐心和一台性能不错的机器）我们基于 PyTorch 和 MLIR 打通了 LLaMA 端到端的推理通路，但是尚未进行完整的性能分析和优化。
目前，我们只应用了针对矩阵乘法的向量化优化，以及针对循环的并行计算优化。
优化和调优的策略仍在开发阶段，因此目前的性能还处于较低水平。我们的初步目标并不是追求极致的性能，而是建立一个标准的端到端通路，在此基础上再做各层级性能优化和调优。
技术路线
技术路线的标准化不仅是我们努力追求的核心原则，而且我们坚信这能够吸引更多的贡献者和合作伙伴，同时还能够有效降低用户的学习门槛和后续的维护成本。
如图所示，我们选择 PyTorch 生态对接各种 AI 模型，选择 MLIR 生态作为 Buddy Compiler 的中间表示。
我们将 Buddy Compiler 作为 Torch Dynamo 的自定义编译器组成整个编译栈，从而希望实现将各种 AI 模型映射到多种硬件架构上的愿景。以下是一些设计点：
使用TorchDynamo作为Trace工具对接AI模型
TorchDynamo 使用 CPython 的 Frame Evaluation API 特性能够做到更加准确的 PyTorch 图的捕捉，详情可以参阅 PyTorch 的文档[5]。除此之外，PyTorch 2.x 提供了全面的编译支持，尤其是提供了非常友好的自定义编译器对接方案[6]。因此，PyTorch 2.x TorchDynamo 成为了我们对接 AI 模型的不二选择。
选择Aten IR作为对接层级
根据 PyTorch 的文档[5]，Core Aten IR 是服务于对接后端的算子集合，它相比 Prime IR 抽象级别也更高。我们倾向于使用较高的抽象级别映射到 MLIR 从而有更多的空间和信息来进行多层优化，因此我们选择 Aten IR 作为对接层级，可以很方便地映射到 Linalg Dialect 和 TOSA Dialect 的各种操作。
使用MLIR Python Bindings实现Dynamo Compiler生成TOSA/Linalg Ops
Buddy Compiler 前端中的 Dynamo Compiler（或者叫做 Dynamo Importer）的作用是将 PyTorch 的 Aten IR 转换到 MLIR，同时也负责将模型参数进行处理和打包。
Dynamo Compiler 遍历从 TorchDynamo 传入的 FX Graph，并且针对每个 Aten IR 节点进行转换，转换过程使用 MLIR Python Bindings 生成目标 MLIR 代码。
Dynamo Compiler 支持对接 Dialect 的优先级设定，即可以选择 Linalg Dialect 或者 TOSA Dialect 作为首选项进行转换。
最终，Dynamo Compiler 负责输出转换后的 MLIR Module 和权重参数。
使用Buddy Compiler工具链进行优化和下降
我们的整个编译通路目前并没有完全使用 Python 脚本完成，而是使用 CMake 将前、中、后端集成起来。
这一定程度上简化了并解耦了前、中、后端的开发流程。编译优化和下降的所有 Pass 注册到 buddy-opt 工具中，它包括所有的上游 Pass 和 Buddy Compiler 中的优化 Pass，因此 buddy-opt 是上游 mlir-opt 的超集。面向通用硬件的编译流程我们使用 MLIR Core Dialect 进行实现，从而达成最大化的复用，我们的工具也因此和所有 LLVM/MLIR 的工具兼容，例如 mlir-translate, llc 等等。
目前我们针对循环采用并行计算的优化，其中相关中间表示和优化 Pass 完全来自上游的 OMP Dialect，可以直接复用并带来不错的优化效果。从此也可以看出统一生态的优势。此外，我们针对粗颗粒度的 Operations 设计了向量化算法进行编译优化，使用 Vector Dialect 也可以实现跨 SIMD/Vector 平台的效果。
如果希望面向特定加速器（例如 Gemmini Accelerator）生成代码，也可以使用 buddy-translate 和 buddy-llc生成到特定于硬件的 LLVM IR，从而最终生成加速器的硬件指令。
Forward函数搭配Buddy Compiler Text Container完成端到端推理
在完成编译优化和下降之后，模型的 Forward 函数将会被构建为共享库。
由于我们很多场景需要在模拟器和开发平台上测试，因此我们没有选择将执行的流程交给 Python 生态，而是进行 AOT 编译，生成可执行文件。
为了配合从 MLIR 构建出来的 Forward 函数实现端到端的推理，我们提供了 C++ 版本的 Text Container 和 MemRef Container 来作为文本输入输出的 Tokenizer 和数据容器。
最终，在 C++ 的 main 函数中将输入的文本和权重参数加载到数据容器，然后调用 Forward 函数进行推理，输出的 Token 再交由 Text Container 进行后处理即可得到最终的文本。
未来工作
我们目前打通了 LLaMA 到 CPU SIMD/Vector 平台的通路，使用 X86 AVX512 进行初步的测试。用于 Vector Dialect 的跨平台性，Arm Neon 和 RISC-V Vector Extesion 也是可以天然支持的，我们正在进行广泛测试。
同时我们也在支持尝试将 Buddy Compiler 中的 Gemmini 加速器支持对接到大模型推理的通路上。
此外，GPU 的优化 Pass 也在开发中。在模型层面，我们希望测试并对接更多的多模态大模型，进一步提升 Buddy Dynamo Compiler 前端的覆盖程度。在完成上述工作后，下一个里程碑是将多模态大模型编译到多种硬件平台上。
总的来说，接下来前、中、后端将会相对解耦地进行开发，最终进行对接。
- 前端：对接更多模型完善算子覆盖程度。
- 中端：进行更详细的性能分析，面向各种计算负载开发不同层级的优化。
- 后端：对 CPU 各 SIMD/Vector 平台进行测试，完成面向 GPU / Gemmini 加速器的端到端通路。
三个部分均相对完备的时候就考虑用 Python 接口将所有工具链包装起来，形成一个更为优雅的使用流程。
总结
如今大模型推理的软件栈也层出不穷，技术路线也各不相同。我们使用的 Torch 2.x + Buddy Compiler 的编译栈设计策略实际上是希望融合 PyTorch 和 MLIR 两大生态进行 AI 计算。当然，在 MLIR 生态里面做大模型推理，我们认为 Torch-MLIR + IREE 目前是相对比较完备的解决方案。nod.ai 的 SHARK Turbine[7] 就使用了这种技术路线。
相比于 Torch-MLIR 搭配 IREE 的组合，Buddy Compiler 更强调 Simple But Powerful 设计，采用极致复用策略和完全代码生成策略。相比于 Torch-MLIR 的 Torch Dialect 层级，Buddy Compiler 更偏向直接复用 TOSA/Linalg 对接 Aten IR；相比于 IREE 覆盖一切后端的 Runtime 和 Execution Engine 的设计，Buddy Compiler 更偏向进行完全代码生成。
Shark Turbine 封面[8]是一个飞机涡轮发动机，这和他们的技术路线非常契合，TorchDynamo + Torch-MLIR + IREE 是一个极其精密且重型的编译栈，这样的"发动机"理论上可以带着他们飞跃任何高山沟壑。相比而言，Buddy Compiler 更像是电动汽车的三电平台，可以以此为基础打造各种性格的电动汽车。对我们来说，LLaMA 的支持不是起点也不是终点，是在探索路上偶遇的一座高山，我们希望翻过它，看看山那边的世界，尤其是开着自己造的车！
致谢
感谢所有 Buddy Compiler 的贡献者，特别感谢一起努力跑通 LLaMA 的伙伴：zhanghb97，weilinquan，xTayEx，EllisLambda，Lester-1，LAJIidea，SForeKeeper，LHY-24，xlinsist，qingqing12138. 同时感谢 OSPP 组委会提供的开源项目席位。
关于Intelligent Computing
Intelligent Computing是由之江实验室和美国科学促进会（AAAS）共同创办，是《科学》合作期刊框架中智能计算领域的第一本开放获取（Open Access）国际期刊。期刊以"面向智能的计算、智能驱动的计算"以及"智能、数据与计算驱动的科学发现"为主题，主要刊载原创研究论文、综述论文和观点论文。
参考链接：