Runtime系统
TVM 支持多种编程语言用于编译器堆栈的开发和部署。解释了 TVM runtime的关键元素。
满足以下需求:
- 部署:从 python/javascript/c++ 语言调用编译函数。
- 调试:在 python 中定义一个函数并从编译的函数中调用。
- 链接:编写驱动程序代码,调用设备特定代码(CUDA),从编译的主机函数中调用。
- 原型:从 python 定义一个 IR pass ,并用 C++ 后端调用。
- Expose:用c++开发的编译器栈到前端(即python)
- 实验:将编译后的函数发送到嵌入式设备,直接在那里运行。
用任何语言定义一个函数,用另一种语言调用。runtime核心尽可能少,以便部署到嵌入式设备。
PackedFunc
PackedFunc是一个简单而优雅的解决方案,可以解决列出的挑战。单个PackedFunc对象代表一个函数调用,调用和被调用可能使用不同的语言。
以下代码块提供了一个 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。需要两个参数:args代表输入参数,rv代表返回值。该函数是type-erased,表明函数签名不限制pass传入或返回的输入类型。调用 PackedFunc 时,将输入参数打包到堆栈上的 TVMArgs,并通过 TVMRetValue 取回结果。
使用 C++ 中的模板技巧,可以像调用普通函数一样调用 PackedFunc。由于type-erased的性质,可以从像 python动态语言中调用 PackedFunc,而无需为创建的每个新类型函数添加额外的胶水代码。以下示例在 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结构。限制了可以传递的可能类型列表。以下是常见的:
- 整数、浮点数和字符串
- PackedFunc 本身
- 编译模块
- 用于张量对象交换的 DLtensor*
- TVM表示 任何IR 中的对象
该限制使实现变得简单,无需序列化。尽管是最小的,但 PackedFunc 足以满足深度学习部署的用例,因为大多数函数只采用 DLtensor 或数字。
由于一个 PackedFunc 可以将另一个 PackedFunc 作为参数,可以将函数从 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 的所有编译器传递函数都作为 PackedFunc 暴露给前端
- 编译后的模块也将编译后的函数返回为 PackedFunc
为了保持runtime最少,将 IR 对象支持与部署runtime隔离。结果runtime需要大约 200K - 600K,具体取决于包含多少runtime驱动程序模块(例如,CUDA)。
与普通函数相比,调用 PackedFunc 的开销很小,只在堆栈中保存了几个值。所以只要不包装小函数就可以了。总之,PackedFunc 是 TVM 中的通用粘合剂,广泛使用支持编译器和部署。
模块
由于 TVM 支持多种类型的设备,需要支持不同类型的驱动程序。必须使用驱动程序 API 来加载内核,以打包格式设置参数并执行内核启动。需要修改驱动程序 API,以便公开的函数是线程安全的。经常需要在 C++ 中实现这些驱动粘合,暴露给用户。不能对每种类型的函数都这样做,所以PackedFunc正确。
TVM 将编译对象定义为Module。用户可以从 Module 中获取编译后的函数作为 PackedFunc。生成的编译代码可以在runtime从 Module 动态获取函数。缓存函数句柄在第一次调用中,并在后续重用。将设备代码和从生成的代码链接到任何 PackedFunc(例如,python)回调。
ModuleNode 是一个抽象类,可以由每种类型的设备实现。到目前为止,支持 CUDA、Metal、OpenCL 和加载动态共享库的模块。这种抽象使得引入新设备变得容易,不需要为每种类型的设备重新生成主机代码。
远程部署
PackedFunc 和 Module 系统,可以轻松地将函数直接传送到远程设备。有一个 RPCModule 序列化参数数据移动,并在远程启动计算。
RPC 服务器本身是最小的,可以捆绑到runtime中。可以在 iPhone/android/raspberry pi 甚至浏览器上启动一个最小的 TVM RPC 服务器。服务器上的交叉编译和测试模块的交付可以在同一个脚本中完成。查看 交叉编译和 RPC以获取更多详细信息。
这种即时反馈带来了很多好处。例如,为了在 iPhone 上测试生成代码的正确性,不再需要从头开始在 swift/objective-c 中编写测试用例——可以使用 RPC 在 iPhone 上执行,将结果复制回来并在主机上进行验证通过 numpy。还可以使用相同的脚本分析。
TVM 对象和编译器堆栈
在 PackedFunc runtime系统构建编译器堆栈 API。为了研究的需要,面临着编译器 API 的不断变化。每当想测试新的原语时,都需要一个新的语言对象或 IR 节点。但是,不希望每次更改 API。除此之外,还希望
- 能够序列化任何语言对象和 IR
- 能够以前端语言探索、打印和操作 IR 对象,进行快速原型设计。
引入了一个名为Object的基类来解决这个问题。编译器堆栈中的所有语言对象都是Object. 每个对象都包含一个字符串 type_key,唯一标识对象的类型。选择 string 而不是 int 作为类型键,因此,Object可以以分散的方式添加新类,而无需将代码添加回中央存储库。为了加快调度速度,在runtime为每个 type_key 分配一个整数 type_index。
通常Object可以在语言中的多个地方引用,使用 shared_ptr 来跟踪引用。使用ObjectRefclass 来表示对Object。 可以粗略地将ObjectRef类看作Object容器的shared_ptr 。还可以定义子类ObjectRef来保存Object。每个Object子类都需要定义 VisitAttr 函数。
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) {}
// ...
};
每个Object子类将override 覆盖,以访问其成员。这是 TensorNode 的一个示例实现。
class TensorNode : public Object {
public:
/*! rief The shape of the tensor */
Array<Expr> shape;
/*! rief data type in the content of the tensor */
Type dtype;
/*! rief the source operation, can be None */
Operation op;
/*! rief the output index from source operation */
int value_index{0};
/*! rief 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);
}
};
在上面的例子中,Operation和Array<Expr>都是 ObjectRef。VisitAttrs 为提供了一个API反射,访问对象的每个成员。可以使用这个函数来访问节点并递归地序列化任何语言对象。允许在前端语言中轻松获取对象的成员。例如,在下面的代码中,访问了 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)
Object可以在不更改前端runtime的情况下,添加到 C++,从而可以轻松地对编译器堆栈进行扩展。这不是将成员暴露给前端语言的最快方法,但可能是最简单的方法之一。使用 Python 进行测试和原型设计,仍然使用 C++ 来完成繁重的工作。
实施细则
PackedFunc 中的每个参数都包含一个union融合值TVMValue 和一个类型代码。这种设计允许动态类型语言直接转换为相应的类型,而静态类型语言在转换过程中进行runtime类型检查。
相关文件是
- C++ API 的packed_func.h
- c_runtime_api.cc用于 C API 以及如何提供回调。
为了支持扩展类型,使用了一个注册系统来注册类型相关的信息,比如在 C++ 中支持 any,更多细节参见扩展类型。