TVM apps extension示例擴展庫
此文件夾包含TVM的示例擴展庫。演示了其它庫如何在C++和Python API中擴展TVM。
該庫擴展了TVM的功能。
python模塊加載新的共享庫，可以使用TVM的python API進行插值。
https://github.com/apache/tvm/tree/main/apps/extension

test_ext.py修改了一些代碼

運行結果

TVM Runtime System
TVM支持多種編程語言用于編譯器堆棧的開發(fā)和部署。將解釋TVMRuntime的關鍵元素。

需要滿足許多有趣的要求：

部署：從python/javascript/c++語言調用編譯后的函數(shù)。

調試：在python中定義一個函數(shù)，從編譯后的函數(shù)調用該函數(shù)。

鏈接：編寫驅動程序代碼，調用設備特定代碼（CUDA），從編譯后的主機函數(shù)中調用。

原型：從Python定義IR傳遞，從C++后端調用。

暴露：C++開發(fā)的編譯器棧到前端（即，Python）

實驗：將編譯后的函數(shù)發(fā)送到嵌入式設備，直接運行。

希望能夠從任何語言定義函數(shù)，從另一種語言調用函數(shù)。希望Runtime核心最小，部署到嵌入式設備。

PackedFunc

PackedFunc是一個簡單但優(yōu)雅的解決方案，可以解決列出的挑戰(zhàn)。單個PackedFunc對象，表示調用者和被調用者可能使用不同語言的函數(shù)調用。

下面的代碼塊提供了C++中的一個例子
#include <tvm/runtime/packed_func.h>

void MyAdd(TVMArgs args, TVMRetValue* rv) {
// automatically convert arguments to desired type.
int a = args[0];
int b = args[1];
// 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。有兩個參數(shù)：args表示輸入?yún)?shù)，rv表示返回值。函數(shù)被類型擦除，這意味著函數(shù)簽名不限制要傳入的輸入類型或要返回的輸入類型。在后端，當調用PackedFunc時，將輸入?yún)?shù)打包到堆棧上的TVMArgs，通過TVMRetValue返回結果。

由于C++中的模板技巧，可以調用PACKEDFUNC，就像普通函數(shù)一樣。由于類型擦除特性，可以從動態(tài)語言（如python）調用PackedFunc，無需為創(chuàng)建的每個新類型函數(shù)，添加額外的粘合代碼。下面的示例在C++中注冊PackedFunc，從Python調用。
// register a global packed function in c++
TVM_REGISTER_GLOBAL(“myadd”)
.set_body(MyAdd);
import tvm

myadd = tvm.get_global_func(“myadd”)

prints 3

print(myadd(1, 2))
PackedFunc的大部分魔力在于TVMArgs和TVMRetValue結構。限制可以傳遞的可能類型的列表。以下是常見的：
? int, float and string
? PackedFunc itself
? Module for compiled modules
? DLTensor* for tensor object exchange
? TVM Object to represent any object in IR
該限制使得實現(xiàn)簡單，無需序列化。盡管PackedFunc是最小的，對于深度學習部署的用例已經足夠了，因為大多數(shù)函數(shù)只接受DLTensor或數(shù)字。
由于一個PACKEDFUNC可以采取另一個PACKEDFUNC作為參數(shù)，所以可以將函數(shù)從Python（如PackedFunc）傳遞到C++。
TVM_REGISTER_GLOBAL(“callhello”)
.set_body([](TVMArgs args, TVMRetValue* rv) {
PackedFunc f = args[0];
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，它允許將PackedFunc嵌入到任何語言中。除了python，支持java和javascript。嵌入式API的這個概念很像Lua，只是沒有新的語言，但是使用C++。
PackedFunc用于編譯器和部署堆棧。
TVM的所有編譯器pass函數(shù)，都作為PackedFunc公開給前端
編譯后的模塊將編譯后的函數(shù)作為PackedFunc返回
為了保持Runtime最小值，將IR對象支持與部署Runtime隔離。根據(jù)包含多少Runtime驅動程序模塊（如CUDA），生成的Runtime大約需要200K-600K。

與普通函數(shù)相比，調用PackedFunc的開銷很小，因為只在堆棧上保存了一些值。因此，只要不包裝小函數(shù)就可以了。總之，PackedFunc是TVM中的通用粘合劑，在TVM中廣泛使用，支持編譯器和部署。
Module
由于TVM支持多種類型的設備，需要支持不同類型的驅動程序。必須使用驅動程序API加載內核，以壓縮格式設置參數(shù)，執(zhí)行內核啟動。需要修補驅動程序API，以便公開的函數(shù)是線程安全的。因此，經常需要在C++中實現(xiàn)這些驅動程序GLUE，公開給用戶。當然不能對每種類型的函數(shù)都這樣做，所以PackedFunc也是答案。
TVM將編譯對象定義為模塊。用戶可以從模塊中以PackedFunc的形式獲取編譯后的函數(shù)。生成的編譯代碼可以在Runtime從模塊中動態(tài)獲取函數(shù)。在第一次調用中緩存函數(shù)句柄，在后續(xù)調用中重用。將設備代碼和回調從生成的代碼鏈接到任何PackedFunc（如python）中。
ModuleNode是一個抽象類，可以由每種類型的設備實現(xiàn)。到目前為止，支持CUDA、Metal、OpenCL和加載動態(tài)共享庫的模塊。這種抽象使得引入新設備變得容易，不需要為每種類型的設備重新生成主機代碼。
遠程部署
PackedFunc和模塊系統(tǒng)可以輕松地將功能直接發(fā)送到遠程設備。在后端，有一個RPCModule，序列化參數(shù)以進行數(shù)據(jù)移動，在遠程啟動計算。

RPC服務器本身是最小的，可以捆綁到Runtime中。可以在iPhone/android/raspberry pi甚至瀏覽器上，啟動最小的TVM RPC服務器。服務器上的交叉編譯和用于測試的模塊的發(fā)布，可以在同一個腳本中完成。有關更多詳細信息，請參考交叉編譯和RPC。
這種即時反饋給了很多好處。例如，為了在iPhone上測試生成代碼的正確性，不再需要從頭開始用swift/objective-c編寫測試用例——可以使用RPC在iPhone上執(zhí)行，將結果復制，通過numpy在主機上進行驗證。可以使用相同的腳本進行分析。
TVM對象和編譯器堆棧
如前所述，在PackedFuncRuntime系統(tǒng)之上，構建編譯器堆棧API。為了研究的需要，面臨著編譯器API的不斷變化。每當想要測試新的原語時，都需要一個新的語言對象或IR節(jié)點。但是，不希望不時更改API。除此之外，
能夠序列化任何語言對象和IRs
能夠使用前端語言探索、打印和操作IR對象，進行快速原型制作。
引入了一個名為Object的基類，解決這個問題。編譯器堆棧中的所有語言對象都是object的子類。每個對象都包含一個字符串類型type_key，該鍵唯一標識對象的類型。選擇string而不是int作為類型鍵，可以以分散的方式添加新的對象類，無需將代碼添加回中心repo。為了簡化調度速度，在Runtime為每個type_key分配一個整數(shù)type_index。

因為通常一個對象可以在語言中的多個位置引用，所以使用一個共享的ptr跟蹤引用。使用ObjectRef類表示對對象的引用。可以粗略地將ObjectRef類視為對象容器的shared_ptr。可以定義子類ObjectRef，保存對象的每個子類型。對象的每個子類都需要定義VisitAttr函數(shù)。
class AttrVisitor {
public:
virtual void Visit(const char* key, double* value) = 0;
virtual void Visit(const char* key, int64_t* value) = 0;
virtual void Visit(const char* key, uint64_t* value) = 0;
virtual void Visit(const char* key, int* value) = 0;
virtual void Visit(const char* key, bool* value) = 0;
virtual void Visit(const char* key, std::string* value) = 0;
virtual void Visit(const char* key, void** value) = 0;
virtual void Visit(const char* key, Type* value) = 0;
virtual void Visit(const char* key, ObjectRef* value) = 0;
// …
};

class BaseAttrsNode : public Object {
public:
virtual void VisitAttrs(AttrVisitor* v) {}
// …
};
每個對象子類都將覆蓋此項訪問成員。下面是TensorNode的一個示例實現(xiàn)。
class TensorNode : public Object {
public:
/*! \brief The shape of the tensor /
Array shape;
/! \brief data type in the content of the tensor /
Type dtype;
/! \brief the source operation, can be None /
Operation op;
/! \brief the output index from source operation /
int value_index{0};
/! \brief constructor */
TensorNode() {}

void VisitAttrs(AttrVisitor* v) final {
v->Visit(“shape”, &shape);
v->Visit(“dtype”, &dtype);
v->Visit(“op”, &op);
v->Visit(“value_index”, &value_index);
}
};
在上面的示例中，操作和數(shù)組都是ObjectRef。VisitAttrs提供了一個反射API，訪問對象的每個成員。可以使用此函數(shù)訪問節(jié)點，遞歸序列化任何語言對象。允許用前端語言輕松獲取對象的成員。例如，在下面的代碼中，訪問了TensorNode的op字段。
import tvm
from tvm import te

x = te.placeholder((3,4), name=“x”)

access the op field of TensorNode

print(x.op.name)
可以在不改變前端運行時的情況下，將新的對象添加到C++，便于對編譯器堆棧進行擴展。這不是向成員公開前端語言的最快方法，但可能是最簡單的方法之一。主要使用Python進行測試和原型開發(fā)，仍然使用C++完成繁重的工作。
實施詳情
PackedFunc中的每個參數(shù)，都包含一個聯(lián)合值TVMValue和一個類型代碼。這種設計允許動態(tài)類型的語言，直接轉換為相應的類型，靜態(tài)類型的語言在轉換過程中，進行運行時類型檢查。
有關檔案如下：
C++ API的PACKEDFUNC.H
c_runtime_api.cc用于c API，如何提供回調。
為了支持擴展類型，使用注冊表系統(tǒng)登記類型相關信息，如C++中的任何支持，參閱擴展類型獲取更多細節(jié)：https://github.com/apache/tvm/tree/main/apps/extension。

參考鏈接：
https://github.com/apache/tvm/tree/main/apps/extension
https://tvm.apache.org/docs/arch/runtime.html

總結

以上是生活随笔為你收集整理的TVM apps extension示例扩展库的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。