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）國際期刊。期刊以 " 面向智能的計算、智能驅動的計算 " 以及 " 智能、數據與計算驅動的科學發現 " 爲主題，主要刊載原創研究論文、綜述論文和觀點論文。
* 本文系量子位獲授權刊載，觀點僅爲作者所有。
— 完 —
量子位 QbitAI
վ ' ᴗ ' ի 追蹤 AI 技術和産品新動态
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
