第一篇剖開一個模型,看到一張 49 個運算的圖。第二篇把那張圖切開,把其中幾塊丟給 NPU。
這篇談的是那個接縫。當 framework 說「這個 subgraph 交給你的硬體跑」,那個呼叫實際上長什麼樣?答案是 Dispatch API。就算你這輩子不會去實作一個,它也值得讀,因為它是設計這層抽象的人親手列出的、一個 NPU 整合必須做對的所有事情。而這剛好就是一份評估晶片的檢查清單。
基於工作上的一個問題我開始讀它,本來只是想確認一件事,結果讀完發現這是拿到那份清單最快的路,比看任何一份廠商簡報都快。
為什麼 TFLite Delegate 不夠用
LiteRT 的文件對「Dispatch API 取代了什麼」講得異常直白。以下轉述自 DISPATCH_API.md 的對照表:
| TFLite Delegate | LiteRT Dispatch | |
|---|---|---|
| 硬體 buffer | 透過 TensorBufferHandle,建立方式各家不同、未標準化 | TensorBuffer |
| Buffer requirement 協商 | 沒有 | TensorBufferRequirements |
| ABI 穩定 | 靠 OpaqueDelegate,但多數 client 用的 C++ 介面並不穩定 | 是 |
| JIT / AOT | delegate 自己搞 JIT,AOT 未標準化 | 兩者皆有,且編譯經由 CompilerPlugin 標準化 |
| 非同步執行 | 沒有 | 支援 |
把「未標準化」和「沒有」那一欄由上往下讀完,動機就很明顯了。在 Delegate 模式下,四個難題,硬體 buffer、buffer 協商、提前編譯、非同步,每一家 vendor 各自解決,或者根本不解決。Dispatch API 不是 API 換皮,它是把這四件事一起拉上來收成一份介面。
ABI 那一列值得多說一句。Dispatch API 全部是 C,這是刻意的。Vendor 的程式碼由 vendor 編譯、放在 vendor 分區出貨,必須能跨 framework 更新繼續運作,而 C++ 介面給不了這個保證。所以用 C。而 LiteRtDispatchInitialize 是在執行期,透過環境選項傳進來的路徑(kLiteRtEnvOptionTagDispatchLibraryDir)找到 vendor library 並動態載入。
Vendor 實際上要實作什麼
Vendor 要填的頂層結構:
typedef struct LiteRtDispatchApi {
LiteRtApiVersion version;
LiteRtDispatchInterface* interface; // 必要
LiteRtDispatchAsyncInterface* async_interface; // 非同步執行
LiteRtDispatchGraphInterface* graph_interface; // 串接多個 executable
LiteRtCustomTensorBufferHandlersDef* tensor_buffer_handlers_def;
} LiteRtDispatchApi;
只有第一個 interface 是必填。其他三個才是有意思的地方,因為 vendor 有沒有填,是真實的能力差異,不是形式問題:
- 沒有
async_interface,每次推理都是阻塞呼叫 - 沒有
graph_interface,沒辦法在裝置上串接多個已編譯的 executable - 沒有 custom buffer handler,你只能用 LiteRT 已知的那幾種 buffer 型別
建立模型時的呼叫序列:
Initialize -> CheckRuntimeCompatibility -> GetCapabilities
-> DeviceContextCreate -> InvocationContextCreate
-> GetInputRequirements / GetOutputRequirements
每次推理:
RegisterTensorBuffer -> AttachInput/AttachOutput -> Invoke -> DetachInput/DetachOutput
目前開源樹裡有四家 vendor 實作:google_tensor、intel_openvino、mediatek、qualcomm。
Requirement 協商才是勝負關鍵
GetInputRequirements 和 GetOutputRequirements 回傳 TensorBufferRequirements,內容指定 buffer 型別、大小、stride、對齊。Runtime 再據此配置滿足條件的 buffer。
第二篇講的「每一刀都有代價」在這裡變得具體。NPU 不接受「一個 float 陣列」,它接受特定型別、特定對齊、特定排列的 buffer。你手上的東西對不上,就會有東西幫你轉,而那次轉換就是一次拷貝。
所以要問 vendor 的問題不是「你的 NPU 支不支援 zero-copy」,因為每一家都會說支援。要問的是:
你的 NPU 接受哪些 buffer 型別?它們和我們相機、GPU 階段產出的型別有交集嗎?
如果答案是「只吃 AHardwareBuffer」而你的管線是 GL texture,那你會永遠每一幀都在轉換,加速器的吞吐量補不回這個。
一份評估晶片的檢查清單
這個系列的起點是工作上一個很實際的問題:要讓端側推理真的成立,而不是變成投影片上的一個項目符號,我們到底需要晶片廠給什麼?
這顆 NPU 是通用可程式化的,還是只服務廠商自己的功能? 這是第一個該問、也最常被跳過的問題。不少電視和行動 SoC 標榜的「AI 引擎」,存在目的是驅動廠商自家的畫質管線,從來沒有以可程式化 runtime 的形式對外開放。規格表上掛在一塊固定功能電路上的 TOPS 數字,對你的價值精確地等於零。
我在音訊 DSP 上看過一模一樣的劇本:device node 在、ls /dev/ 看得到、名字裡甚至就寫著 DSP,然後你花了兩天才搞清楚它永遠不會跑你的模型。
有沒有 runtime,而且它在裝置上嗎? 要具體的 .so 名稱、它出貨在哪個分區、版本是多少。Dispatch API 是動態載入 vendor library 的,那個 library 不在 image 裡,就沒有東西可以載入。
有沒有 compiler,而且它能進 CI 嗎? 這跟 runtime 是兩個問題,但很容易搞混,我自己一開始就混在一起問。裝置上有 runtime,不代表你做得出東西給它跑。如果廠商的模型編譯器需要自己一份授權,那會限制你整條 release pipeline。而且照第二篇說的,如果廠商只支援 AOT,compiler 就坐在關鍵路徑上,沒有 JIT 可以退。
支援哪些量化格式、哪些 op? op 覆蓋窄就會切出碎片圖,也就是第二篇那個比 CPU 還慢的死法。
接受哪些 buffer 型別? 同上。
NPU 需不需要專屬或連續記憶體?要多少? 在記憶體吃緊的裝置上,這是直接跟其他所有東西搶預算。
userspace 是 64-bit 嗎? 很容易忘記問。在 32-bit ARM 上你會失去 int8 dot product 指令,而那正是量化模型在 CPU 上跑得快的主因,所以連你的退路都是慢的。
老實說,「幾 TOPS」根本不該出現在這份清單上。吞吐量當然重要,但上面每一個問題都能獨立地讓 TOPS 這個數字變得無關緊要,所以先把這些問完,把吞吐量留到最後當比較用的那一項。
這三篇下來我的收穫
寫完這三篇,改變最大的是我對「難點在哪裡」的認知。我原本以為端側推理主要是模型架構和量化的問題,結果反而大部分是介面的問題:晶片承認自己支援哪些 op、接受哪些 buffer 型別、compiler 你有沒有資格執行、runtime 在不在 image 裡。
如果你現在正在為某個端側功能評估晶片,vendor 的標頭檔會比 vendor 的簡報回答你更多問題。我會從那裡開始看。
資料來源:google-ai-edge/LiteRT commit dc32e93 的 DISPATCH_API.md、COMPILER_PLUGIN.md、JIT_COMPILATION.md,以及 litert/vendors/ 底下的各家實作。全部 Apache 2.0,一個下午讀得完。