真真夜夜的炼金工坊

感謝 up 主 ZOMI 醬：https://space.bilibili.com/517221395

為什麼需要 AI 編譯器#

面臨的問題#

傳統編譯器與 AI 編譯器#

傳統編譯器與 AI 編譯器主要有兩大差異：

1. IR 差異：AI 編譯器的 IR 與傳統編譯器的 IR 所抽象出來的概念和意義並不相同。

   1. AI 編譯器一般會有 high-level IR，用來抽象描述深度學習模型中的運算，如：Convolution、Matmul 等，甚至部分會有 Transformer 帶有圖的結構。
   2. 傳統編譯器相對而言 low-level IR，用於描述基本指令運算，如 load、store 等。有了 high-level IR，AI 編譯器在描述深度學習模型類 DSL 更加方便。

2. 優化策略：AI 編譯器面向 AI 領域，優化時引入更多領域特定知識，從而進行更 high-level，更加 aggressive 優化手段。如：

   1. AI 編譯器在 high-level IR 執行算子融合，傳統編譯器執行類似 loop fusion 時候，往往更加保守。缺點是可能會導致調試執行信息跟蹤難；
   2. AI 編譯器可以降低計算精度，比如 int8、fp16、bf16 等，因為深度學習對計算精度不那麼敏感。但傳統編譯器一般不執行改變變量類型和精度等優化。

AI 編譯器的發展階段#

推理場景：輸入 AI 框架訓練出來的模型文件，輸出能夠在不同硬體高效執行的程序；
訓練場景：輸入高級語言表示的神經網絡代碼，輸出能夠在不同硬體高效執行的程序；

1. 什麼是訓練場景？什麼是推理場景嗎？
2. 搞 AI 編譯器為什麼要了解算法呢？
3. 搞 AI 算子為什麼要了解編譯器？

什麼是 AI 編譯器#

1. Python 為主的動態解釋器語言前端
2. 多層 IR 設計，包括圖編譯、算子編譯、代碼生成
3. 面向神經網絡、深度學習的特定優化
4. DSA 芯片架構的支持

AI 編譯器的發展階段#

Stage I： 樸素的 AI 編譯器#

TensorFlow 早期版本，基於神經網絡的編程模型，主要進行了 graph 圖和 ops 算子兩層抽象。

• 圖層：通過聲明式的編程方式，以靜態圖方式執行，執行前進行硬體無關和硬體相關的編譯優化。硬體無關的優化，如表達式化簡、常量折疊、自動微分等；硬體相關的優化包括算子融合、內存分配等。
• 算子層：通常採用手寫 kernel 的方式，如在 NVIDIA GPU 上基於 CUDA kernel 實現大量的 .cu 算子或者依賴於 CuDNN 算子優化庫。

在表達上：

• 靜態圖的表達式非 Python 原生，開發者主要通過框架提供 Python API 顯示構圖，易用性上不好；

在性能上：

• DSA 專用加速芯片出現加劇了性能上的挑戰；
• 算子層提供的算子粒度和邊界提前確定後，無法充分發揮硬體的性能；
• 硬體廠商的提供的算子優化庫也未必最優
  • 1）模型和 shape 確定情況下，可能還有更優算子實現；
  • 2）在 SIMT 和 SIMD 架構中，Scheduling、Tilling 都有有很大的空間。

Stage II 專用的 AI 編譯器#

在表達上：

• PyTorch 靈活表達 API 方式成為 AI 框架參考標杆，圖層的神經網絡編譯器主要就是考慮如何把類 PyTorch 的表達轉換到圖層的 IR 進行優化。
• 類 PyTorch 的 Python 原生表達，靜態化轉換；
• AI 專用編譯器架構，打開圖算邊界進行融合優化；

在性能上：

• 打開計算圖和算子的邊界，進行重新組合優化，發揮芯片的算力。計算圖層下發子圖中的算子打開成小算子，基於小算子組成的子圖，進行編譯優化，包括 buffer fusion、水平融合等，關鍵是大算子怎樣打開、小算子如何重新融等。

• 表達分離：計算圖層和算子層仍然分開，算法工程師主要關注圖層的表達，算子表達和實現主要是框架開發者和芯片廠商提供。
• 功能泛化：對靈活表達上的動靜態圖轉換、動態 Shape、稀疏計算、分佈式並行優化等複雜的需求難以滿足。
• 平衡效率和性能：算子實現上在 Schedule、Tiling、Codegen 上缺乏自動化手段，門檻高，開發者既要了解算子計算邏輯，又要熟悉硬體體系架構。

Stage III 通用 AI 編譯器#

• 圖算統一表達，實現融合優化
• 算子實現上自動 Schedule、Tiling、Codegen，降低開發門檻
• 更泛化優化能力，實現動静統一、動態 Shape、稀疏性、高階微分、自動並行等
• 包括編譯器、運行時，異構計算、邊緣到數據中心都模塊化表示和組合，並專注於可用性

AI 編譯器的通用架構#

IR 中間表達#

編譯器主要分為前後端，分別針對於硬體無關和硬體相關的處理。每一個部分都有自己的 IR，每個部分也會對進行優化：

• High-level IR：用於表示計算圖，其出現主要是為了解決傳統編譯器中難以表達深度學習模型中的複雜運算這一問題，為了實現更高效的優化所以新設計了一套 IR。
• Low-level IR：能夠在更細粒度的層面上表示模型，從而能夠針對於硬體進行優化，文中將其分為三類。

Frontend 前端優化#

構造計算圖後，前端將應用圖級優化。因為圖提供了計算全局概述，所以更容易在圖級發現和執行許多優化。前端優化與硬體無關，這意味著可以將計算圖優化應用於各種後端目標。前端優化分為三類：

1. 節點級優化，如 Zero-dim-tensor elimination、Nop Elimination
2. 塊級優化，如代數簡化、常量折疊、算子融合
3. 數據流級優化，如 Common sub-expression elimination、DCE

Backend 後端優化#

• 特定硬體的優化
  • 目標針對特定硬體體系結構獲取高性能代碼。1）低級 IR 轉換為 LLVM IR，利用 LLVM 基礎結構生成優化的 CPU/GPU 代碼。2）使用領域知識定制優化，這可以更有效地利用目標硬體。
• 自動調整
  • 由於在特定硬體優化中用於參數調整的搜索空間巨大，因此有必要利用自動調整來確定最佳參數設置。1）Halide/TVM 允許調度和計算表達分開，使用自動調節來得出較佳配置。2）應用多面體模型 Polyhedral model 進行參數調整。
• 優化內核庫
  • 廠商特定優化內核庫，廣泛用於各種硬體上的加速 DL 訓練和推理。特定優化原語可以滿足計算要求時，使用優化的內核庫可顯著提高性能，否則可能會受到進一步優化的約束。

AI 編譯器的挑戰與思考#

XLA：優化機器學習編譯器#

XLA（加速線性代數）是一種針對特定領域的線性代數編譯器，能夠加快 TensorFlow 模型的運行速度，而且可能完全不需要更改源代碼。

XLA：圖層下發子圖中的算子打開成小算子，基於小算子組成的子圖進行編譯優化，包括 buffer fusion、水平融合等，關鍵是大算子怎樣打開、小算子如何重新融合、新的大算子如何生成，整體設計主要通過 HLO/LLO/LLVM IR 實現，所有 Pass 規則都是手工提前指定。

TVM：端到端深度學習編譯器#

為了使得各種硬體後端的計算圖層級和算子層級優化成為可能，TVM 從現有框架中取得 DL 程序的高層級表示，並產生多硬體平台後端上低層級的優化代碼，其目標是展示與人工調優的競爭力。

TVM：分為 Relay 和 TVM 兩層，Relay 關注圖層，TVM 關注算子層，拿到前端子圖進行優化，Relay 關注算子間融合、TVM 關注新算子和 kernel 生成，區別在於 TVM 開放架構，Relay 目標是可以接入各種前端，TVM 也是一個可以獨立使用的算子開發和編譯的工具，算子實現方面採用 Compute（設計計算邏輯）和 Schedule（指定調度優化邏輯）分離方案。

TENSOR COMPREHENSIONS：神經網絡語言編譯器#

TensorComprehensions 是一種可以構建 just in time (JIT) 系統的語言，程序員可通過該語言用高級編程語言去高效的實現 GPU 等底層代碼。

TC：算子計算邏輯的較容易實現，但 Schedule 開發難，既要了解算法邏輯又要熟悉硬體體系架構，此外圖算邊界打開後，小算子融合後，會生成新算子和 kernel，新算子 Schedule 很難生成，傳統方法定義 Schedule 模板；TC 希望通過 Polyhedral model 實現 auto schedule。

nGraph：兼容所有框架的深度學習系統編譯器#

nGraph 的運作應作為深度學習框架當中更為複雜的 DNN 運算的基石，並確保在推理計算與訓練及推理計算之間取得理想的效率平衡點。

挑戰#

1. 動態 Shape 和動態計算圖：

   • 現狀：主流 AI 編譯器主要針對特定靜態 shape 輸入進行編譯優化，對包含控制流語義的動態計算圖支持有限。
   • 問題：AI 應用場景中存在大量需要動態計算圖的任務需求。儘管前端可以嘗試將計算圖改寫為靜態計算圖，或將適合編譯的部分子圖展開以供優化，但這些方法不能解決所有問題。
   • 實例：某些 AI 任務（如金字塔結構的檢測模型）無法通過人工改寫來靜態化，編譯器在這些情況下難以有效優化。

2. Python 編譯靜態化：

   方法	描述
   Python JIT 虛擬機	如 PyPy 或 CPython，試圖在 Python 解釋執行的基礎上增加 JIT 編譯加速能力，但兼容性存在問題。
   修飾符方式	如 torch.jit.script、torch.jit.trace、torch.fx、LazyTensor、TorchDynamo。儘管這些手段已經廣泛應用，但仍缺乏一個完備的方案。

   Python JIT 虛擬機

**修飾符方式**

- **AI編譯器在 Python 靜態化方面的挑戰**：
    - **類型推導**：從Python的動態類型到編譯器IR的靜態類型。
    - **控制流表達**：如if、else、while、for等控制語句的表達。
    - **靈活的語法和數據類型轉換**：針對Slice、Dict、Index等操作的處理。
    - **JIT編譯性能**：無論是Tracing Based還是AST Transform，都需要額外的編譯開銷。

3. 發揮硬體性能，特別是 DSA 類芯片：

• 現狀：DSA（專用硬體架構）在 AI 訓練和推理中的重要性日益增加，如 CPU 的 SIMD 單元、GPU 的 SIMT 架構、華為昇騰 Ascend 的 Cube 核等。

• 挑戰：

  • 性能優化依賴圖算融合：需要圖層和算子層獨立優化，無法充分發揮芯片性能，需要圖算融合優化，如子圖切分、子圖內垂直融合優化和水平並行優化。
  • 優化複雜度提升：涉及標量、向量、張量和加速指令的多級存儲結構，導致 Kernel 實現的調度、分塊、向量化、張量化複雜。

• 解決方案：

  • 打開圖和算子的邊界進行重新組合優化，以充分發揮芯片性能。
  • 融合多種優化手段：垂直融合優化（如 Buffer Fusion）和水平並行優化（如 Data Parallel）。
  • 自動生成重新組合優化後的子圖的 Kernel 代碼，如調度、分塊、向量化。

4. 處理神經網絡特性：自動微分、自動並行等：

5. 易用性與性能兼顧

其他問題#

• 圖算能否統一表達，統一編譯優化，形成通用的 AI 編譯器？
  當前的 AI 框架下，圖層和算子層是分開表達和優化，算法工程師主要是接觸圖層表達，AI 框架、芯片、kernel 開發工程師主要是負責算子的表達，未來能否會有 IR 打破圖算之間的 GAP，未來在 AI + 科學計算、AI + 大數據等新場景驅動下，使得計算圖層和算子層不再清晰，能否統一 AI 的編譯優化？

• 完全的自動並行是否可行？
  自動並行能根據用戶輸入的串行網絡模型和提供的集群資源信息自動進行分佈式訓練，通過採用統一分佈式計算圖和統一資源圖設計可支持任意並略和各類硬體集群資源上分佈式訓練，並且還能利用基於全局代價模型的規劃器自適應為訓練任務選擇硬體感知的並行策略。實際上自動並行是一個策略搜索問題，策略搜索能夠在有限的搜索空間找到一個次優的答案，但是真正意義上的自動並行能否做到需要進一步思考和驗證。

• AI 芯片需要編譯器嗎？AI 芯片需要 AI 編譯器嗎？
  AI 芯片對於編譯器的依賴取決於芯片本身的設計。越靈活的芯片對於編譯器的依賴會越大。在 AI 芯片設計之初，有 CISC 風格把優化在芯片內部解決。但是隨著專用領域的演化，為了支持更加靈活的需求，AI 芯片本身會在保留張量指令集和特殊內存結構的前提下越來越靈活。未來的架構師需要芯片和系統協同設計，自動化也會越來越多地被應用到專用芯片中去。

未來#

• 編譯器形態：分開推理和訓練，AOT 和 JIT 兩種編譯方式並存
• IR 形態：需要有類似於 MLIR 對 AI 統一的 IR 表達
• 自動並行：提供跨機器、跨節點自動並行的編譯優化能力
• 自動微分：提供高階微分的計算方式，並方便對圖進行操作
• kernel 自動生成：降低開發門檻，快速實現高效和高泛化的算子