【TVM教程】TVM 運行時系統詳情 - 機器學習,人工智能,深度學習,openai,知識超神經HyperAI 博客

TVM 現已更新到 0.21.0 版本，TVM 中文文檔已經和新版本對齊。

Apache TVM 是一個深度的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →Apache TVM

TVM 支持多種編程語言用於編譯器棧的開發和部署。在本説明中，我們將解釋 TVM 運行時的關鍵組成部分。

VM 的運行時系統需要滿足多種看似相互矛盾但又非常關鍵的需求：

部署（Deployment）：能夠在 Python / JavaScript / C++ 等語言中調用已編譯的函數。
調試（Debug）：允許用户在 Python 中定義函數，並從已編譯的代碼中反向調用。
鏈接（Linking）：需要編寫驅動端代碼來調用設備端實現（如 CUDA kernel），並且運行時需要能從主機端代碼中調用它們。
原型開發（Prototyping）：支持在 Python 中創建 IR Pass，並能從 C++ 後端調用。
接口暴露（Frontend Exposure）：編譯器的核心邏輯由 C++ 實現，但必須便捷地暴露給 Python 等前端語言。
實驗與部署（Experiment & Deployment）：能夠將編譯好的函數直接傳輸並運行在嵌入式設備上。

我們希望能夠在任何語言中定義函數並在另一種語言中調用。我們還希望運行時核心儘可能小，以便部署到嵌入式設備上。

PackedFunc

PackedFunc是我們找到的一個簡單但優雅的解決方案來解決列出的挑戰。一個 PackedFunc 對象就表示一次函數調用，而調用方和被調用方可以處於不同的語言環境中。

下面的代碼塊提供了一個 C++ 示例

#include <tvm/ffi/function.h>

void MyAdd(ffi::PackedArgs args, ffi::Any* rv) {
  // automatically convert arguments to desired type.
  int a = args[0].cast<int>();
  int b = args[1].cast<int>();
  // automatically assign value return to rv
  *rv = a + b;
}

void CallPacked() {
  PackedFunc myadd = PackedFunc(MyAdd);
  // get back 3
  int c = myadd(1, 2);
}

在上面的代碼塊中，我們定義了一個 PackedFunc MyAdd。它接受兩個參數：args 表示輸入參數，rv 表示返回值。該函數是類型擦除的，這意味着函數簽名不會限制傳入或返回值的類型。在底層，當我們調用一個 PackedFunc 時，它會將輸入參數打包成 ffi::PackedArgs 放在棧上，並通過 ffi::Any 獲取返回結果。

得益於 C++ 中的模板機制，我們可以像調用普通函數一樣調用 PackedFunc。由於其類型擦除的特性，我們可以在諸如 Python 這樣的動態語言中調用 PackedFunc，而不需要為每一種新函數類型額外編寫 glue 代碼。下面的例子展示瞭如何在 C++ 中註冊一個 PackedFunc，並在 Python 中調用它。

// register a global packed function in c++
TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def_packed("myadd", MyAdd);
}

import tvm

myadd = tvm.get_global_func("myadd")
# prints 3
print(myadd(1, 2))

PackedFunc 的大部分「魔力」來自 ffi::PackedArgs 和 ffi::Any 這兩個結構。我們對可傳遞的類型做了限制，常見的類型包括：

int、float 和 string
PackedFunc 本身
Module，用於表示已編譯模塊
DLTensor*，用於張量對象交換
TVM Object，用於表示 IR 中的任意對象

這種限制使得實現變得簡單，無需序列化。即使實現精簡，PackedFunc 在深度學習部署的場景中依然綽綽有餘，因為大多數函數只需要處理 DLTensor 或數字。

由於一個 PackedFunc 可以將另一個 PackedFunc 作為參數傳遞，因此我們可以將 Python 中的函數（轉換為 PackedFunc）傳遞給 C++。

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def_packed("callhello", [](ffi::PackedArgs args, ffi::Any* rv) {
    ffi::Function f = args[0].cast<ffi::Function>();
    f("hello world");
  });
}

import tvm

def callback(msg):
  print(msg)

# convert to PackedFunc
f = tvm.convert(callback)
callhello = tvm.get_global_func("callhello")
# prints hello world
callhello(f)

TVM 提供了一個最小化的 C API minimum C API，它允許我們將 PackedFunc 嵌入到任意語言中。除了 Python 以外，目前還支持 java 和 javascript。這種嵌入式 API 的設計理念與 Lua 很相似，只不過我們並沒有創造一門新的語言，而是直接使用了 C++。

關於 PackedFunc 有一個有趣的事實：我們在編譯器棧和部署棧中都使用它。

TVM 中所有編譯器 Pass 函數都以 PackedFunc 的形式暴露給前端
已編譯模塊同樣以 PackedFunc 的形式返回已生成的函數

為了保持運行時儘可能精簡，我們將 IR Object 支持從部署運行時中分離開來。最終生成的運行時大小大約為 200K - 600K，具體取決於包含的運行時驅動模塊數量（例如 CUDA）。

調用 PackedFunc 相比普通函數的開銷很小，只多做了一些棧上值保存。因此，只要不頻繁包裝非常小的函數，這樣的開銷是可以接受的。總的來説，PackedFunc 是 TVM 的通用“膠水層”，我們在編譯和部署模塊中都大量依賴它。

組件

由於 TVM 支持多種不同類型的硬件設備，我們也需要支持對應的不同驅動程序。我們必須使用這些驅動 API 來加載內核、以打包形式設置參數並啓動內核執行。同時，我們還需要對驅動 API 進行封裝，以確保暴露給用户的接口是線程安全的。因此，我們通常會在 C++ 中編寫這些驅動層 Glue 代碼，並通過 PackedFunc 將其暴露給用户。顯然，我們不可能為每類函數都單獨編寫接口，因此 PackedFunc 再次成為解決方案。

TVM 將編譯結果抽象為一個 Module。

用户可以從 Module 中以 PackedFunc 的形式獲取已編譯函數。生成的代碼在運行時可以動態地從 Module 中獲取目標函數，並在第一次調用時緩存句柄，後續複用。這使得我們可以在生成代碼中鏈接設備端函數，並調用任意 PackedFunc（例如 Python 回調）。

ModuleNode 是一個抽象類，不同設備類型可以各自實現。例如，我們已支持 CUDA、Metal、OpenCL 以及動態庫（Shared Library）。這種抽象設計使得引入新設備變得簡單，而無需重新生成每種設備的主機端代碼。

遠程部署

PackedFunc 和 Module 系統也使得我們可以將函數直接部署到遠程設備上。在底層，我們提供了一個 RPCModule，它負責序列化參數、進行數據傳輸，並在遠程設備上啓動計算。

RPC 服務器本身非常精簡，可以直接與運行時一起打包。我們可以在 iPhone、Android、樹莓派甚至瀏覽器中啓動一個最小化的 TVM RPC 服務器。交叉編譯、模塊打包與測試都可以在同一個腳本中完成。更多細節可參考 tutorial-cross-compilation-and-rpc。

這種即時反饋帶來了顯著優勢。例如，當我們希望驗證生成的代碼在 iPhone 上的正確性時，不再需要手動用 Swift/Objective-C 重寫測試樣例——我們可以直接使用 RPC 在 iPhone 上執行代碼，將結果複製回主機，並使用 numpy 進行驗證。同樣，我們也可以使用同一個腳本進行性能分析。

TVM 對象與編譯器棧

如前所述，我們在 PackedFunc 運行時系統之上構建了編譯器棧的 API。由於研究需求，編譯器 API 經常需要不斷變化。當我們想要測試新的語言原語時，就需要引入新的語言對象或 IR 節點。但是我們又不希望頻繁修改 API。此外，我們還希望：

能夠序列化任意語言對象和 IR；
能夠在前端語言中探索、打印和操作 IR 對象，以便進行快速原型開發。

為了解決這些問題，我們引入了一個基類Object。編譯器棧中的所有語言對象都是 Object 的子類。每個對象都包含一個字符串 type\_key，用於唯一標識對象類型。我們選擇字符串而不是整數作為類型鍵的原因是：這樣可以以去中心化方式添加新的 Object 類，而無需往中心倉庫中添加代碼。為了加速調度，我們會在運行時為每個 type\_key 分配一個整數 type\_index。

由於一個 Object 通常會在語言中被多個地方引用，我們使用 shared\_ptr 來管理對象引用。ObjectRef 類用於表示對 Object 的引用，可以將其視為指向Object容器的 shared\_ptr。我們也可以定義 ObjectRef 的子類來對應不同的 Object子類型。每個 Object 子類都需要實現 RegisterReflection 函數。

每個Object子類會重寫該函數來註冊其成員。下面是 IntImmNode 的示例實現：

class IntImmNode : public PrimExprNode {
public:
  /*! \brief the Internal value. */
  int64_t value;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, PrimExprNode);
};
// in cc file
TVM_FFI_STATIC_INIT_BLOCK() { IntImmNode::RegisterReflection(); }

RegisterReflection為我們提供了一個反射接口，用於註冊對象的成員。我們可以利用這個函數遞歸地訪問並序列化任何語言對象。同時，它也使我們可以在前端語言中輕鬆訪問對象的字段。例如：

import tvm

x = tvm.tir.IntImm("int32", 1)
# access the value field of IntImmNode
print(x.value)

新的 Object 可以僅在 C++ 中添加而無需修改前端運行時，從而方便擴展編譯器棧。需要注意的是，這種機制不是訪問成員的最高性能方式，但它是最簡單的方法之一。我們發現這種方式非常適合我們的目的：用 Python 進行測試和原型開發，而真正的計算和重工作交由 C++ 完成。

實現細節

PackedFunc 中的每個參數由一個聯合體 TVMValue 和一個類型碼組成。這樣的設計使得動態類型語言可以直接轉換到對應類型，而靜態類型語言則可以在轉換過程中執行運行時類型檢查。

相關文件包括：

packed\_func.h —— C++ API
c\_runtime\_api.cc —— C API 以及如何提供回調支持

為了支持擴展類型，我們使用了一個註冊表系統來註冊類型相關信息，例如允許 C++ 中對 any的支持。更多詳情可參考：Extension types。

與運行時相關的信息

Vulkan Runtime

超神經HyperAI 博客

超神經HyperAI 博客

博客 / 詳情