Buddy Compiler 端到端 LLaMA2-7B 推理示例已經(jīng)合并到 buddy-mlir倉(cāng)庫(kù)[1]主線。

我們?cè)?Buddy Compiler 的前端部分實(shí)現(xiàn)了面向 TorchDynamo 的第三方編譯器，從而結(jié)合了 MLIR 和 PyTorch 的編譯生態(tài)。

目前，前端部分可以覆蓋 LLaMA 計(jì)算圖，轉(zhuǎn)換到 MLIR 后我們集成了部分向量化和并行優(yōu)化，并在 AVX512 平臺(tái)上進(jìn)行了測(cè)試。整個(gè)推理過(guò)程可以跑通但還需要大量?jī)?yōu)化。以下是相關(guān)鏈接和現(xiàn)狀：

• [E2E] Buddy Compiler 端到端 LLaMA2-7B 推理示例[2]
• [E2E] 上述端到端推理示例目的是展示編譯棧設(shè)計(jì)，并非完備的 LLaMA 問(wèn)答工具
• [Frontend] Buddy Dynamo Compiler[3]
• [Midend] 集成面向矩陣乘法的向量化以及面向循環(huán)的并行優(yōu)化
• [Backend] 端到端示例在 X86 AVX512 機(jī)器上進(jìn)行測(cè)試（Ubuntu 22.04）
• [WIP] 開(kāi)發(fā)并集成各種優(yōu)化（現(xiàn)在速度太慢）
• [WIP] 在多種 CPU 后端進(jìn)行測(cè)試（Arm Neon, RVV, etc.）
• [WIP] 對(duì)接多種硬件平臺(tái)（GPU, Gemmini Accelerator, etc.）
• [WIP] 增加前端覆蓋程度（Stable Diffusion, CLIP, Whisper, etc.）

概述

AI 大模型的爆發(fā)為軟硬件設(shè)計(jì)帶來(lái)了新的抓手和機(jī)會(huì)。隨著模型的規(guī)模、類(lèi)型、模態(tài)的增加和發(fā)散，單一軟硬件技術(shù)棧覆蓋各種場(chǎng)景的能力越來(lái)越受限。這種趨勢(shì)也加深了我對(duì)軟硬件生態(tài)重要性的認(rèn)識(shí)。在我看來(lái)，整個(gè)生態(tài)的最重要的三條設(shè)計(jì)原則如下（重要程度由高到低）：

• 技術(shù)路線標(biāo)準(zhǔn)化（生態(tài)的立足之本）
• 上手門(mén)檻低（足夠多的貢獻(xiàn)者和玩家是生態(tài)繁榮的關(guān)鍵）
• 優(yōu)化上限高（決定了生態(tài)的發(fā)展?jié)摿Γ?/li>

我們基于以上思考打造了 Buddy Compiler， 致力于實(shí)現(xiàn)軟硬件協(xié)同設(shè)計(jì)生態(tài)。我們的目標(biāo)是實(shí)現(xiàn)從領(lǐng)域特定編程語(yǔ)言（DSL）到領(lǐng)域特定硬件架構(gòu)（DSA）的編譯流程和協(xié)同設(shè)計(jì)。

本文將介紹如何使用 Buddy Compiler 完成 LLaMA 2 的端到端推理。同時(shí)，我們也會(huì)分享 Buddy Compiler 的整體設(shè)計(jì)理念和未來(lái)的規(guī)劃。具體的構(gòu)建流程請(qǐng)參閱此處的文檔[4]（大模型的構(gòu)建需要十足的耐心和一臺(tái)性能不錯(cuò)的機(jī)器）我們基于 PyTorch 和 MLIR 打通了 LLaMA 端到端的推理通路，但是尚未進(jìn)行完整的性能分析和優(yōu)化。

目前，我們只應(yīng)用了針對(duì)矩陣乘法的向量化優(yōu)化，以及針對(duì)循環(huán)的并行計(jì)算優(yōu)化。

優(yōu)化和調(diào)優(yōu)的策略仍在開(kāi)發(fā)階段，因此目前的性能還處于較低水平。我們的初步目標(biāo)并不是追求極致的性能，而是建立一個(gè)標(biāo)準(zhǔn)的端到端通路，在此基礎(chǔ)上再做各層級(jí)性能優(yōu)化和調(diào)優(yōu)。

技術(shù)路線

技術(shù)路線的標(biāo)準(zhǔn)化不僅是我們努力追求的核心原則，而且我們堅(jiān)信這能夠吸引更多的貢獻(xiàn)者和合作伙伴，同時(shí)還能夠有效降低用戶(hù)的學(xué)習(xí)門(mén)檻和后續(xù)的維護(hù)成本。

如圖所示，我們選擇 PyTorch 生態(tài)對(duì)接各種 AI 模型，選擇 MLIR 生態(tài)作為 Buddy Compiler 的中間表示。

我們將 Buddy Compiler 作為 Torch Dynamo 的自定義編譯器組成整個(gè)編譯棧，從而希望實(shí)現(xiàn)將各種 AI 模型映射到多種硬件架構(gòu)上的愿景。以下是一些設(shè)計(jì)點(diǎn)：

使用TorchDynamo作為T(mén)race工具對(duì)接AI模型

TorchDynamo 使用 CPython 的 Frame Evaluation API 特性能夠做到更加準(zhǔn)確的 PyTorch 圖的捕捉，詳情可以參閱 PyTorch 的文檔[5]。除此之外，PyTorch 2.x 提供了全面的編譯支持，尤其是提供了非常友好的自定義編譯器對(duì)接方案[6]。因此，PyTorch 2.x TorchDynamo 成為了我們對(duì)接 AI 模型的不二選擇。

選擇Aten IR作為對(duì)接層級(jí)

根據(jù) PyTorch 的文檔[5]，Core Aten IR 是服務(wù)于對(duì)接后端的算子集合，它相比 Prime IR 抽象級(jí)別也更高。我們傾向于使用較高的抽象級(jí)別映射到 MLIR 從而有更多的空間和信息來(lái)進(jìn)行多層優(yōu)化，因此我們選擇 Aten IR 作為對(duì)接層級(jí)，可以很方便地映射到 Linalg Dialect 和 TOSA Dialect 的各種操作。

使用MLIR Python Bindings實(shí)現(xiàn)Dynamo Compiler生成TOSA/Linalg Ops

Buddy Compiler 前端中的 Dynamo Compiler（或者叫做 Dynamo Importer）的作用是將 PyTorch 的 Aten IR 轉(zhuǎn)換到 MLIR，同時(shí)也負(fù)責(zé)將模型參數(shù)進(jìn)行處理和打包。

Dynamo Compiler 遍歷從 TorchDynamo 傳入的 FX Graph，并且針對(duì)每個(gè) Aten IR 節(jié)點(diǎn)進(jìn)行轉(zhuǎn)換，轉(zhuǎn)換過(guò)程使用 MLIR Python Bindings 生成目標(biāo) MLIR 代碼。

Dynamo Compiler 支持對(duì)接 Dialect 的優(yōu)先級(jí)設(shè)定，即可以選擇 Linalg Dialect 或者 TOSA Dialect 作為首選項(xiàng)進(jìn)行轉(zhuǎn)換。

最終，Dynamo Compiler 負(fù)責(zé)輸出轉(zhuǎn)換后的 MLIR Module 和權(quán)重參數(shù)。

使用Buddy Compiler工具鏈進(jìn)行優(yōu)化和下降

我們的整個(gè)編譯通路目前并沒(méi)有完全使用 Python 腳本完成，而是使用 CMake 將前、中、后端集成起來(lái)。

這一定程度上簡(jiǎn)化了并解耦了前、中、后端的開(kāi)發(fā)流程。編譯優(yōu)化和下降的所有 Pass 注冊(cè)到 buddy-opt 工具中，它包括所有的上游 Pass 和 Buddy Compiler 中的優(yōu)化 Pass，因此 buddy-opt 是上游 mlir-opt 的超集。面向通用硬件的編譯流程我們使用 MLIR Core Dialect 進(jìn)行實(shí)現(xiàn)，從而達(dá)成最大化的復(fù)用，我們的工具也因此和所有 LLVM/MLIR 的工具兼容，例如 mlir-translate, llc 等等。

目前我們針對(duì)循環(huán)采用并行計(jì)算的優(yōu)化，其中相關(guān)中間表示和優(yōu)化 Pass 完全來(lái)自上游的 OMP Dialect，可以直接復(fù)用并帶來(lái)不錯(cuò)的優(yōu)化效果。從此也可以看出統(tǒng)一生態(tài)的優(yōu)勢(shì)。此外，我們針對(duì)粗顆粒度的 Operations 設(shè)計(jì)了向量化算法進(jìn)行編譯優(yōu)化，使用 Vector Dialect 也可以實(shí)現(xiàn)跨 SIMD/Vector 平臺(tái)的效果。

如果希望面向特定加速器（例如 Gemmini Accelerator）生成代碼，也可以使用 buddy-translate 和 buddy-llc生成到特定于硬件的 LLVM IR，從而最終生成加速器的硬件指令。

Forward函數(shù)搭配Buddy Compiler Text Container完成端到端推理

在完成編譯優(yōu)化和下降之后，模型的 Forward 函數(shù)將會(huì)被構(gòu)建為共享庫(kù)。

由于我們很多場(chǎng)景需要在模擬器和開(kāi)發(fā)平臺(tái)上測(cè)試，因此我們沒(méi)有選擇將執(zhí)行的流程交給 Python 生態(tài)，而是進(jìn)行 AOT 編譯，生成可執(zhí)行文件。

為了配合從 MLIR 構(gòu)建出來(lái)的 Forward 函數(shù)實(shí)現(xiàn)端到端的推理，我們提供了 C++ 版本的 Text Container 和 MemRef Container 來(lái)作為文本輸入輸出的 Tokenizer 和數(shù)據(jù)容器。

最終，在 C++ 的 main 函數(shù)中將輸入的文本和權(quán)重參數(shù)加載到數(shù)據(jù)容器，然后調(diào)用 Forward 函數(shù)進(jìn)行推理，輸出的 Token 再交由 Text Container 進(jìn)行后處理即可得到最終的文本。

未來(lái)工作

我們目前打通了 LLaMA 到 CPU SIMD/Vector 平臺(tái)的通路，使用 X86 AVX512 進(jìn)行初步的測(cè)試。用于 Vector Dialect 的跨平臺(tái)性，Arm Neon 和 RISC-V Vector Extesion 也是可以天然支持的，我們正在進(jìn)行廣泛測(cè)試。

同時(shí)我們也在支持嘗試將 Buddy Compiler 中的 Gemmini 加速器支持對(duì)接到大模型推理的通路上。

此外，GPU 的優(yōu)化 Pass 也在開(kāi)發(fā)中。在模型層面，我們希望測(cè)試并對(duì)接更多的多模態(tài)大模型，進(jìn)一步提升 Buddy Dynamo Compiler 前端的覆蓋程度。在完成上述工作后，下一個(gè)里程碑是將多模態(tài)大模型編譯到多種硬件平臺(tái)上。

總的來(lái)說(shuō)，接下來(lái)前、中、后端將會(huì)相對(duì)解耦地進(jìn)行開(kāi)發(fā)，最終進(jìn)行對(duì)接。

• 前端：對(duì)接更多模型完善算子覆蓋程度。
• 中端：進(jìn)行更詳細(xì)的性能分析，面向各種計(jì)算負(fù)載開(kāi)發(fā)不同層級(jí)的優(yōu)化。
• 后端：對(duì) CPU 各 SIMD/Vector 平臺(tái)進(jìn)行測(cè)試，完成面向 GPU / Gemmini 加速器的端到端通路。

三個(gè)部分均相對(duì)完備的時(shí)候就考慮用 Python 接口將所有工具鏈包裝起來(lái)，形成一個(gè)更為優(yōu)雅的使用流程。

總結(jié)

如今大模型推理的軟件棧也層出不窮，技術(shù)路線也各不相同。我們使用的 Torch 2.x + Buddy Compiler 的編譯棧設(shè)計(jì)策略實(shí)際上是希望融合 PyTorch 和 MLIR 兩大生態(tài)進(jìn)行 AI 計(jì)算。當(dāng)然，在 MLIR 生態(tài)里面做大模型推理，我們認(rèn)為 Torch-MLIR + IREE 目前是相對(duì)比較完備的解決方案。nod.ai 的 SHARK Turbine[7] 就使用了這種技術(shù)路線。

相比于 Torch-MLIR 搭配 IREE 的組合，Buddy Compiler 更強(qiáng)調(diào) Simple But Powerful 設(shè)計(jì)，采用極致復(fù)用策略和完全代碼生成策略。相比于 Torch-MLIR 的 Torch Dialect 層級(jí)，Buddy Compiler 更偏向直接復(fù)用 TOSA/Linalg 對(duì)接 Aten IR；相比于 IREE 覆蓋一切后端的 Runtime 和 Execution Engine 的設(shè)計(jì)，Buddy Compiler 更偏向進(jìn)行完全代碼生成。

Shark Turbine 封面[8]是一個(gè)飛機(jī)渦輪發(fā)動(dòng)機(jī)，這和他們的技術(shù)路線非常契合，TorchDynamo + Torch-MLIR + IREE 是一個(gè)極其精密且重型的編譯棧，這樣的“發(fā)動(dòng)機(jī)“理論上可以帶著他們飛躍任何高山溝壑。相比而言，Buddy Compiler 更像是電動(dòng)汽車(chē)的三電平臺(tái)，可以以此為基礎(chǔ)打造各種性格的電動(dòng)汽車(chē)。對(duì)我們來(lái)說(shuō)，LLaMA 的支持不是起點(diǎn)也不是終點(diǎn)，是在探索路上偶遇的一座高山，我們希望翻過(guò)它，看看山那邊的世界，尤其是開(kāi)著自己造的車(chē)！

致謝

感謝所有 Buddy Compiler 的貢獻(xiàn)者，特別感謝一起努力跑通 LLaMA 的伙伴：zhanghb97，weilinquan，xTayEx，EllisLambda，Lester-1，LAJIidea，SForeKeeper，LHY-24，xlinsist，qingqing12138. 同時(shí)感謝 OSPP 組委會(huì)提供的開(kāi)源項(xiàng)目席位。