--- title: AI Compiler -- TVM --- # AI Compiler -- TVM [TOC] ## 1. Introduction AI Compiler是一種專門為深度學習模型進行最佳化與部署的工具。 傳統編譯器負責將程式碼轉換為機器可執行的指令,而 AI Compiler 的任務則是將神經網路模型轉換為在特定硬體(如 CPU、GPU、NPU、TPU 或 ASIC)上以最高效能運行的形式。它會針對計算圖進行優化,例如合併運算(operator fusion)、移除冗餘節點、重新排序計算流程,並可搭配量化、模型壓縮等技術,降低運算量與記憶體需求。透過這些過程,AI Compiler 能大幅提升模型推論速度、降低功耗,並改善在行動裝置或嵌入式平台上的執行效率。 在傳統軟體開發中,GCC 或 Clang 將 C/C++ 程式碼編譯成機器碼;同樣地,AI Compiler 將以 PyTorch、TensorFlow 或 ONNX 格式定義的神經網路模型,轉換並優化為可在特定硬體(如 CPU、GPU、或專用 AI 加速器)上高效執行的低階程式碼。因此,AI Compiler 在 AI 模型的實際應用中扮演關鍵角色,也是部署高效能 AI 系統不可或缺的一環。 ## 2. Purpose 在 AI 開發中,會面臨一個著名的 $M \times N$ 難題: - $M$ 種訓練框架:可能用 PyTorch, TensorFlow, Keras, MXNet 等訓練模型。 - $N$ 種硬體平台:想把模型跑在 NVIDIA GPU, Intel CPU, ARM 手機晶片, 甚至是 FPGA 或專用加速器(NPU)上。 AI Compiler 的核心目標是解決模型在不同平台上部署的效能與相容性問題: - **模型轉換與部署(Translation & Deployment)**:將高階模型(如 PyTorch 模型)轉換為中間表示法(IR),再翻譯成目標硬體可執行的格式(如 C 語言原始碼或二進位檔),使其能運行於嵌入式系統或資源受限的環境中。 - **效能優化(Optimization)**:透過自動化的圖形優化技術(如算子融合),減少記憶體存取次數並提升推論速度,無需人工手寫低階組語。 - **跨平台支援(Portability)**:作為一個中介層,讓同一個模型可以被編譯並部署到異質平台(Heterogeneous Platforms)上,例如從伺服器端遷移至邊緣裝置。 ## 3. Architecture AI Compiler 架構可分為以下幾個層級: - 前端(Frontend) - 中間表示法(Intermediate Representation, IR) - 優化層(Optimization Layer) - 後端與代碼生成(Backend & Codegen) - 執行環境(Runtime) 在開源 AI Compiler 框架中,**Apache TVM(Tensor Virtual Machine)** 是目前最重要、最完整、最具彈性的 AI Compiler 框架之一。 TVM 是 AI 領域的 LLVM。它打破了深度學習框架與硬體之間的藩籬,透過自動化編譯技術,讓模型跑得更快、更省資源,並且能運行在更多地方。TVM 提供了一個中間層,可以用任何框架訓練模型,然後將其編譯並優化,最後高效地運行在任何硬體上。 以下將會針對 TVM 的架構與功能進行介紹。 ### 3.1 TVM Compiler #### 3.1.1 Frontend & IR 接收來自 TensorFlow、PyTorch 或 ONNX 的模型輸入,並將模型轉換為計算圖(Computation Graph)。例如 TVM 使用 Relay 作為高階中間表示法(Intermediate Representation, IR),詳細記錄每一層的輸入形狀(Shape)與資料型態(Data Type)。 自製 AI Compiler 多使用現成的 `tvm.relay.frontend` ,除非有使用自定義算子,或輸入不是標準框架而是自製的 Domain Specific Language(DSL)。 #### 3.1.2 Compiler & Optimizer 這是 TVM 最核心的部分,負責改寫計算圖以提升效率: - **算子融合(Operator Fusion)**:將多個連續的運算(例如 `Conv2d + BiasAdd + ReLU`)合併為一個複合節點(Composite Node)。這能大幅減少中間數據在記憶體中的搬運次數。 - **量化(Quantization)**:將浮點數運算轉換為整數運算(如 float32 轉 int8),並插入對應的 Quantize/Dequantize 節點。 - **圖形分割(Partitioning)**:將計算圖切割,標記出哪些部分由特定的硬體加速器(如 DLA)處理,哪些由 CPU 處理。 - **圖層級優化(Graph-level Optimization)**:例如將多個運算(如 `Conv + ReLU`)合併成一個,減少記憶體存取次數。 - **算子層級優化(Operator-level Optimization)**:在更低一層的 TIR (Tensor Intermediate Representation)進行,Relay 最後會 Lowering(降階)成 TIR,再針對具體硬體調整矩陣乘法、迴圈(Loop)的執行方式。 - **AutoTVM / Ansor(Auto-scheduler)**:傳統編譯器依靠人工手寫規則來優化代碼。TVM 使用機器學習來優化機器學習,自動在硬體上進行數千次微小的測試(Tuning),搜尋出該硬體上運算最快的參數組合。AutoTVM 是基於「模板(Template)」的搜尋。需要預先定義好搜尋空間,而 Ansor 則無模板(Template-free)。它更強大,能自動生成搜尋空間,是比 AutoTVM 更先進的技術,能自動生成優化策略而無需人工寫模板。 #### 3.1.3 Types of Relay Pass TVM 的 Relay Pass 生態系非常龐大,總數可能有 50 個以上。它們涵蓋了從最底層的型別推導,到高層的代數簡化,再到記憶體分配的方方面面。以下舉例幾項常用類別。 1. **簡化與清理類(Simplification & Cleanup)** 這類 Pass 負責把模型變乾淨,去除多餘的運算。這通常是優化的第一步。 - `SimplifyInference`(推理簡化) - 作用:去除訓練階段才需要的算子,例如 `Dropout`。同時也會把 `BatchNorm` 融合進前面的 `Conv2d` 或 `Dense` 層(因為在推理時,BN 其實就是一組固定的乘法和加法)。 - 意義:硬體不需要實作 `BatchNorm`,因為它會消失。 - `FoldConstant`(常數摺疊) - 作用:預先計算所有輸入都是常數的節點。例如 `x = 2 + 3` 會直接變成 `x = 5`。這也包括預先計算權重的某些變換(如轉置)。 - 意義:減少執行時期的計算量。 - `DeadCodeElimination`(死碼消除) - 作用:移除那些計算了但沒有被輸出的節點。 2. **算子融合類(Operator Fusion)** 這是深度學習編譯器最重要的優化之一。 - `FuseOps`(算子融合) - 作用:將多個小算子合併成一個大 Kernel。 - 例子:`Conv2d` -> `BiasAdd` -> `ReLU`。如果不融合,GPU/NPU 需要讀寫記憶體 3 次;融合後,只需讀寫 1 次。 - 意義:如果硬體支援「`Conv+ReLU`」一次做完,這個 Pass 能自動把圖切好,Codegen 就能收到一個包含 `Conv` 和 `ReLU` 的複合函數。 3. **數據佈局與轉換類(Layout & Data Transform)** 這類 Pass 負責調整資料在記憶體中的排布,以適應硬體特性。 - `AlterOpLayout`(改變算子佈局) - 作用:自動轉換 Tensor 的形狀格式。例如將標準的 NCHW 轉換為 NHWC,或是更複雜的 Block 格式(如 NCHW16c,專為向量化指令設計)。 - 意義:很多自研加速器喜歡特定的記憶體排列(例如為了對齊寬度),不需要在 Frontend 強制使用者改模型,只需定義好這個 Pass,TVM 會自動插入 `LayoutTransform` 節點來幫忙搬運資料。 - `ToMixedPrecision` / `Quantize`(混合精度 / 量化) - 作用:將 FP32 模型轉換為 FP16 或 INT8。需要準備一小批真實資料(Calibration Dataset)餵給模型跑一次,用來統計每一層 Tensor 的數值範圍(Min/Max),才能決定量化參數(Scale/Zero-point)。若沒有校準,直接量化的模型精度會掉得非常嚴重。 - 意義:如果加速器只跑 INT8,這個 Pass 流程是必經之路。 4. **BYOC 專用類(Hardware Integration)** 如果要用 BYOC(Bring Your Own Codegen)接入自研硬體,以下這幾個 Pass 是標準流程。 - `MergeComposite` - 作用:這是基於「模式匹配(Pattern Matching)」的融合。可以定義一個 Pattern(eg. `Conv2d + Add + Sigmoid`),這個 Pass 就會掃描整個圖,把符合這形狀的地方打包起來。 - 意義:這是 BYOC 的第一步,用來找出硬體能吃的子圖。 - `AnnotateTarget` - 作用:標記哪些節點要給 CPU 跑,哪些要給加速器跑。 - `PartitionGraph` - 作用:根據上面的標記,把 Relay 圖真正切開。被切出來的部分會變成一個獨立的函數,準備丟給 Codegen。 5. **基礎設施類(Infrastructure & Correctness)** 這些 Pass 通常在幕後默默運作,確保 Relay 程式碼是合法的。 - `InferType`(型別推導) - 作用: 算出每個節點的 Output Shape 和 Data Type。 - 注意: 如果自己寫了 Pass 修改了 Graph 結構(例如把一個 Node 換掉了),通常必須馬上跑一次 InferType,否則後面的 Pass 會報錯崩潰。 - `ToANormalForm`(轉換為 ANF) - 作用: 把巢狀的表達式拉平。這對於後續的分析和 Codegen 比較友善。 - `Inline`(內聯) - 作用: 把小的函數直接展開到呼叫它的地方,減少函數呼叫開銷。 6. **記憶體與動態形狀處理(Memory & Dynamic Shapes)** 這對硬體設計者來說是個痛點,因為硬體通常不喜歡動態記憶體。 - `DynamicToStatic`(動態轉靜態) - 作用: 嘗試把動態 Shape(例如 Input 大小不固定)轉換成靜態 Shape。如果硬體只支援固定 Input Size,這個 Pass 是救星。 - `ManifestAlloc`(顯式記憶體分配) - 作用: 這通常在編譯流程的後期。它會把抽象的 Tensor 運算,變成具體的「申請記憶體 -> 計算 -> 釋放記憶體」的指令。 7. **代數與數學優化(Algebraic Simplification)** 這類優化類似傳統編譯器(如 GCC/LLVM)做的事。 - `EliminateCommonSubexpr`(CSE - 公共子表達式消除) - 作用: 如果 `a = x + y` 算了兩次,它會把第二次計算刪掉,直接用第一次的結果。 - 意義: 節省硬體運算週期。 - `FastMath` - 作用: 用精度較低但在硬體上更快的數學函數來替換標準函數(例如用近似的 `exp` 或 `tanh`)。 8. **量化專用 Pass(Quantization / QNN)** 如果做 INT8 加速器,會大量接觸 `relay.qnn.transform` 下的 Pass。 - `Legalize` - 作用: 把某些硬體不支援的複雜 QNN 算子,轉換成標準的 Relay 算子組合。 - `CanonicalizeOps` - 作用: 規範化算子,讓量化邏輯更統一。 ### 3.2 TVM Backend & Codegen #### 3.2.1 Function 這部分主要有以下幾項功能: - **代碼生成(Code Generation)**:將優化後的 IR 翻譯成目標語言。 - **邏輯程式碼**:生成控制流程的 C 語言程式碼(`model.c`),包含函式呼叫與記憶體管理(`malloc/free`)。 - **權重資料**:將模型參數序列化為二進位檔案(`weight.bin`)或 C 陣列(`weight.c`)。 當我們在 TVM 中指定 `target="c"` 時,TVM 會啟用標準的 C Codegen。它的核心任務是將經過優化後的 TIR (Tensor IR)翻譯成標準的 C99 原始碼。 1. 運作原理(Mechanism) - AST Traversal(語法樹遍歷):TVM 的 `CodeGenC` 類別會遍歷 TIR 的語法樹。 - 映射邏輯: - 將 TIR 的 `For` 節點翻譯成 C 語言的 `for (int i=0; i<n; ++i) { ... }`。 - 將 TIR 的 `Load/Store` 節點翻譯成陣列存取 `buffer[index]。 - 將 TIR 的 `IfThenElse` 翻譯成 `if (...) { ... }`。 2. 產出結構(Output Structure)生成的 C 程式碼通常包含兩個主要部分: - PackedFunc Wrapper:這是 TVM Runtime 用來呼叫函數的通用介面。它接收 `void* args`(輸入張量) 和 `void* type_code`(型別資訊),負責將這些通用指標轉型為具體的資料指標。 - Kernel Implementation:真正執行數學運算的邏輯,通常是一層又一層的巢狀迴圈(Nested Loops)。 同時 TVM 也有提供 **Bring Your Own Codegen(BYOC)**,這是一種允許開發者自定義代碼生成邏輯的機制(後述)。 ### 3.3 TVM Runtime TVM Runtime 是 AI 模型在目標裝置(Target Device)上的執行環境。如果 TVM Compiler 是「遊戲開發者」(製作光碟),Runtime 就是「遊戲主機」(讀取光碟並執行遊戲)。 #### 3.3.1 TVM Runtime 特性: - 輕量化(Lightweight): 僅包含載入模型和執行運算的最少代碼(通常僅幾百 KB 到幾 MB)。 - 獨立性: 執行時不需要安裝 PyTorch、TensorFlow 或完整的 TVM Compiler。 - 高可攜性: 核心由 C++ 實作,可運行在手機、樹莓派、微控制器、FPGA SOC 等各種平台。 #### 3.3.2 核心工作流程(Workflow) Runtime 的生命週期通常包含四個標準步驟: - 載入(Load): 讀取編譯好的二進位模型檔案(如 `.so`, `.tar`)。 - 輸入(Set Input): 將原始資料(圖片、語音)轉換為 TVM 專用的 Tensor 格式(NDArray)並填入模型記憶體。 - 執行(Run): - 觸發計算圖的執行。 - 若是 CPU 模型,呼叫 CPU 指令。 - 若是加速器模型,呼叫底層 Driver 發送指令給硬體。 - 輸出(Get Output): 等待運算完成,將結果搬運回 Host 端供使用者讀取。 #### 3.3.3 硬體開發者視角:Runtime 與自研加速器 - 對於開發 AI 加速器的實驗室,Runtime 扮演 Host(CPU)與 Device(加速器)的溝通橋樑。 - 關鍵機制 - PackedFunc: TVM 使用 PackedFunc 機制將 C++ 函數包裝起來,讓 Python 或其他語言可以呼叫。 - Driver 的位置: 實作的硬體驅動程式(Driver Code)會被編譯進 Runtime 中。當上層呼叫 `run()` 時,Runtime 實際上是在執行自己寫的 C/C++ 代碼來控制硬體(例如寫入暫存器、配置 DMA)。 ## 4. Development & Debugging Tools 在 AI Compiler 的開發過程中,如何以視覺的方式看見模型被編譯成什麼樣子,以及如何遠端在嵌入式板子上除錯,是兩個最關鍵的需求。 ### 4.1 TVM RPC(Remote Procedure Call) 對於嵌入式開發者(如開發 ARM 板子或 FPGA SOC)來說,直接在目標裝置(Device)上編譯模型通常是不切實際的,因為板子的算力與記憶體有限。TVM RPC 機制完美解決了這個問題。 - 架構分離(Host-Device Separation): - Host 端(PC/Server):負責耗資源的任務。載入模型、執行編譯(Relay -> Graph -> Codegen)、生成最終的 `tar` 或 `so` 檔。 - Device 端(開發板):只需運行一個輕量的 `tvm_rpc_server`。 - 工作流程: - Host 端編譯模型。 - Host 透過網路(TCP)將編譯好的二進位檔「上傳」到 Device 的 RAM 中。 - Host 發送指令,要求 Device 執行推論(Inference)。 - Device 執行完畢,將結果與執行時間(Profile)回傳給 Host。 - 優勢:開發者無需在板子上安裝龐大的編譯環境(LLVM/GCC/Python),也無需頻繁插拔 SD 卡複製檔案,大幅縮短「Tune-Compile-Deploy」的循環週期。 ### 4.2 Visualization Tools 當編譯器執行了大量的 Pass(如 3.1.3 提到的 `FuseOps` 或 `LayoutTransform`)後,原來的模型結構已經面目全非。這時需要視覺化工具來確認編譯器的行為是否符合預期。 - Netron:AI 領域通用的模型檢視器。 - 除錯應用:在 Relay 優化的不同階段(例如 Optimize 之後),可以將 IR 匯出為 `.json` 或 `.relay` 格式。 - 將檔案丟入 Netron,可以直觀地看到: - 算子融合是否成功:檢查 `Conv2d` 和 `ReLU` 是否真的變成了一個節點。 - Layout 是否改變:檢查 Tensor 是否從 NCHW 變成了 NHWC。 - 屬性檢查:點擊節點可查看 Stride, Padding, Dtype 等詳細資訊。 ### 4.3 Profiling & Debugger - Graph Executor Debugger:TVM 內建的除錯器。 - 它可以吐出每一層(Layer-wise)的執行時間。這對於硬體開發者至關重要,能幫忙找出「哪一層跑得特別慢」,從而針對該算子優化硬體指令或排程。 - 它可以Dump出每一層的中間輸出值(Intermediate Output),用於跟 Golden Data(PyTorch 跑出的結果)比對,抓出精度誤差發生在哪一層。 ## 5. Advanced Quantization: Calibration 在第 3 章提到了量化(Quantization),但從 FP32 轉換到 INT8 並非直接截斷數值,否則精度會大幅下降。這中間必須經過校準(Calibration)流程。 ### 5.1 Function - FP32 的數值範圍很大(10^-38^ ~ 10^38^),而 INT8 只有 -128 到 127。 - 如果模型中的某一層 Activation 值分布在 -10.5 到 20.1 之間,我們需要找出一個最佳的 Scale(縮放比例)和 Zero-point(零點),將這個範圍映射到 INT8 上。 ### 5.2 Calibration Workflow 這通常發生在 Compile 階段的前期。 - 準備校準資料集(Calibration Dataset):從訓練資料中隨機抽取一小部分真實圖片(例如 100 張)。 - 模擬推論:用 FP32 模式跑這 100 張圖,並統計每一層 Layer 的輸出數值分布(Histogram)。 - 計算參數:根據統計結果決定截斷範圍(Threshold)。 - MinMax:直接取最大最小值(對雜訊敏感)。 - KL Divergence(KL 散度):尋找截斷後資訊損失最小的範圍(目前主流做法,精度較高)。 - 生成量化模型:將計算出的 Scale/Zero-point 寫入 IR,並將權重轉換為 INT8。 ## 6. Scope of Customization 在基於 TVM 開發自研硬體編譯器時,通常保留 TVM 的 Frontend 與中層優化,僅替換與硬體緊密相關的後半段。以下是自製 AI Compiler 必須自行開發的三大核心模組: ### 6.1 Custom Pattern & Fusion Rules 3.1 節中提到 TVM Compiler 會將多個算子合併為一個複合節點,這能大幅減少中間數據在記憶體中的搬運次數。TVM 預設知道如何融合 CPU/GPU 的算子(如 Conv2d + ReLU),但它不知道自研硬體的特性與特殊算子,因此需要定義 **Pattern Table**,告訴 TVM 可以將哪些特殊算子的元素打包成單一算子。 例如如果自研加速器硬體支援「`Conv + Bias + ReLU + Pool`」一次做完,需要寫一個 Python 的 Pattern Matcher,讓 TVM 只要在圖上看到 `Conv->Bias->ReLU->Pool` 這種結構,將它們打包成一個名為 `my_lab_fused_op` 的黑盒子函數。 ### 6.2 Custom Codegen(BYOC) TVM 有提供 **Bring Your Own Codegen(BYOC)**,這是一種允許開發者自定義代碼生成邏輯的機制,利用這個功能,可以利用 TVM 框架開發屬於自研硬體加速器的 Codegen。 #### 6.2.1 Function 當 TVM 把圖切好之後,Codegen 需要負責生成各個算子對應的 C 程式碼。需要撰寫 `codegen.py`,這是一個翻譯官,他會遍歷 Relay 的計算圖(AST),生成個算子對應的指令。 翻譯邏輯範例: - 讀到 `nn.conv2d` → 輸出 `dla_set_config(...); dla_start_conv(...)`; - 讀到 `Data Tensor` → 輸出 `dla_malloc(...)`; - 讀到 `Weight` → 將浮點數權重轉為硬體需要的 `Hex` 格式並輸出為陣列。 #### 6.2.2 TVM Codegen vs. Custom Codegen TVM 預設的 C Codegen 產出的是標準 for-loop 運算。它產出的 Code 是給 CPU 跑的運算邏輯,而不是給加速器跑的 控制指令(Control Code/Configuration)。Custom Codegen 不需要生成「乘法運算」,而是生成「設定暫存器(Register Setting)」的代碼。 使用這兩者的差異主要在於 「通用性 vs. 專用性」 以及 「Runtime 環境的限制」。以下針對這幾個關鍵點進行比較: 1. **硬體適配性(Hardware Adaptation)** - 現成的 Codegen(General Purpose): - 通常只支援標準硬體架構(如 x86, ARM CPU, NVIDIA GPU)。 - 生成的程式碼是通用的 C 或 Assembly,無法呼叫自己設計的加速器指令。 - 自研 AI Compiler(BYOC): - 需求:目標硬體是 Eyeriss DLA(Deep Learning Accelerator)或特定的 ASIC。這些硬體有自己獨特的 Driver API(eg. `qconv2d_relu_maxpool_dla(...)`)。 - 差異:通用的 Codegen 不知道硬體的樣子,也不知道怎麼呼叫 Driver。需要自定義編譯器,將 Relay 的高階算子(eg. `nn.conv2d`)精準地映射到 ASIC Driver API 上。 2. **優化層級(Optimization & Fusion)** - 現成的 Codegen: - 算子融合(Fusion)規則是通用的(例如常見的 `Conv+ReLU`)。 - 無法得知硬體有什麼特殊算子或能力。 - 自研 AI Compiler: - 需求:假設 DLA 硬體設計得很強,可以一口氣做完 `Conv + Bias + ReLU + MaxPool` 而不需要把中間資料寫回記憶體。 - 差異:透過 `fuse.py`,可以定義專屬於硬體的融合規則(Pattern)。如果用現成的 Codegen,它可能會把 MaxPool 拆開來做,導致資料要在記憶體來回移動,浪費了硬體設計的優勢。 3. **Runtime 環境與相依性(Runtime Overhead)** - 現成的 Codegen: - Standard TVM 預設生成的 C code,通常需要依賴 TVM Runtime Library(`libtvm_runtime.so`)才能執行。這是一個數 MB 大小的函式庫,負責記憶體管理和函數派發。 - 對於資源極度受限的微控制器(MCU)或嵌入式系統,這個 Runtime 可能太肥大了。 - 自研 AI Compiler: - 需求:作業目標是產生 Bare-metal(裸機)等級的 C code。 - 差異:在 `codegen.py` 裡生成的程式碼是 Pure C,只依賴標準的 `malloc/free` 和自己寫的簡單 `runtime_cpu.c`。這讓程式碼非常輕量,完全不需要掛載龐大的 TVM Runtime 就能運行,這對於嵌入式部署至關重要。 - 註:TVM 官方也有極輕量的 MicroTVM CRT 方案,Runtime 只有幾 KB,專門給 MCU 用,只是自研加速器通常傾向自己寫。 4. **開發與除錯透明度(Debuggability)** - 現成的 Codegen: - 通常像個黑盒子,生成的程式碼非常複雜且難以閱讀(往往是一堆 void* 指標和自動生成的變數名)。 - 自研 AI Compiler: - 需求:需要用 valgrind 分析記憶體,並且要能看懂每一層 layer 到底發生了什麼事。 - 差異:透過自己寫 Codegen,可以控制變數命名、記憶體配置的方式(`malloc`),產出的 `model.c` 結構清晰、可讀性高。這能清楚看到每一層 Layer 的呼叫順序,方便進行效能分析與除錯。 ### 6.3 Device Runtime & Driver Custom Codegen 生成出來的 C code 會呼叫一些底層函數(API),這些函數也必須實作。 - Memory Management:實作 `device_malloc` 和 `device_free`。需要管理加速器上的 SRAM 或專用 DRAM 位址。 - Instruction Issue:實作發送指令的函數。例如透過 AXI Bus 或 PCIe 將 Codegen 算好的 Config 寫入硬體的暫存器。 - Synchronization:實作 `dla_wait()`,讓 CPU 等待硬體跑完(通常透過 Polling 或 Interrupt)。 為什麼不能用 TVM 的?因為 TVM 的 Runtime (`libtvm_runtime`) 是通用的,它無法控制用 Verilog/Chisel 寫出來的特殊硬體介面,且自己寫 Runtime 可以確保最輕量,不會放入不會被使用的函式庫,對於有空間須極端縮小要求的部署非常重要。