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);
}

In the above codeblock, we defined a PackedFunc MyAdd. It takes two arguments : args represents input arguments and rv represents return value. The function is type-erased, which means that the function signature does not restrict which input type to pass in or type to return. Under the hood, when we call a PackedFunc, it packs the input arguments to ffi::PackedArgs on stack, and gets the result back via 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))

Most of the magic of PackedFunc lies in ffi::PackedArgs and ffi::Any structure. We restrict a list of possible types which can be passed. Here are the common ones:

• int, float 和字符串

• 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 提供了 minimum C API，这使得能够将 PackedFunc 嵌入到任何语言中。除了 Python，到目前为止还支持 java 和 javascript。这种嵌入式 API 的理念与 Lua 非常相似，不同之处在于没有创造一种新的语言，而是使用 C++。

关于 PackedFunc 的有趣的事实是，既用它来处理编译器，也用它来处理部署堆栈。

• TVM 中的所有编译器 pass 函数都被作为 PackedFunc 暴露给前端。

• 编译后的模块同样会以 PackedFunc 的形式返回编译好的函数。

为了最小化运行时的开销，将 IR 对象支持从部署运行时中分离出来。这样得到的运行时大约需要 200K 到 600K的空间，具体取决于包含了多少运行时驱动模块（例如 CUDA）。

调用 PackedFunc 与普通函数的开销很小，因为它仅在堆栈上保存了几个值。因此，只要不包装小型函数，就没有问题。总之，PackedFunc 是 TVM 中广泛使用的万能粘合剂(glue)，用它来支持编译器和部署工作。

Module#

由于 TVM 支持多种类型的设备，需要为不同类型的驱动程序提供支持。必须使用驱动 API 来加载内核，以打包格式设置参数并执行内核启动。还需要对驱动 API 进行修补，以确保暴露的函数是线程安全的。因此，通常需要用 C++ 实现这些驱动粘合代码，并向用户公开。显然，不能为每种类型的函数都这样做，因此 PackedFunc 再次成为解决方案。

TVM 将编译后的对象定义为 Module。用户可以从 Module 中获取已编译函数作为 PackedFunc。生成的编译代码可以在运行时动态从 Module 中获取函数。它在首次调用时缓存函数句柄，并在后续调用中重用。使用这种方法将设备代码和回调链接到任何由生成的代码中的 PackedFunc （例如，Python）。

ModuleNode 抽象类，每种类型的设备都可以实现它。到目前为止，支持 CUDA、Metal、OpenCL 以及动态共享库的加载模块。这种抽象使得引入新设备变得简单，并且无需针对每种设备类型重新生成主机代码。

远程部署#

PackedFunc 和 Module 系统也使得将函数直接部署到远程设备变得简单。在底层，有 RPCModule，它序列化参数以实现数据迁移，并在远程启动计算过程。

RPC 服务器本身是最小化的，并且可以打包到运行时环境中。可以在 iPhone、Android、树莓派甚至浏览器上启动最小化的 TVM RPC服务器。服务器上的交叉编译和模块的测试发布可以在相同的脚本中完成。有关更多详细信息，请参阅：tutorial-cross-compilation-and-rpc。

这种即时反馈给我们带来了许多优势。例如，为了测试在 iPhone 上生成代码的正确性，不再需要从头开始编写 Swift/Objective-C 的测试用例——可以使用 RPC 在 iPhone 上执行，将结果复制回主机并通过 NumPy 进行验证。还可以使用相同的脚本进行性能分析。

TVM Object 与编译器堆栈#

正如之前提到的，在 PackedFunc 运行时系统之上构建了编译器堆栈 API。由于研究需求，面临编译器 API 的持续变化。每当希望测试新的基元(primitives)时，就需要新的语言对象或 IR 节点。然而，不希望 API 频繁变动。此外，我们还希望

• 能够将任何语言对象和 IRs 序列化。

• 能够在前端探索、打印和操作 IR 对象，使用前端语言进行快速原型设计。

为了解决这个问题，引入了基类 Object 。编译器堆栈中的所有语言对象都是 Object 的子类。每个对象都包含字符串类型的 type_key，这个键唯一标识了对象的类型。选择使用字符串而非整数作为类型键，这样新的 Object 类就可以以去中心化的方式添加，而无需将代码重新提交到中央仓库。为了提高调度速度，在运行时为每个 type_key 分配整数类型的 type_index。

Since usually one Object could be referenced in multiple places in the language, we use a shared_ptr to keep track of reference. We use ObjectRef class to represent a reference to the Object. We can roughly view ObjectRef class as shared_ptr to the Object container. We can also define subclass ObjectRef to hold each subtypes of Object. Each subclass of Object needs to define the RegisterReflection function.

Each Object subclass will override this to register its members. Here is an example implementation of 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);
  }

  static constexpr const char* _type_key = "ir.IntImm";
  TVM_DECLARE_FINAL_OBJECT_INFO(IntImmNode, PrimExprNode);
};
// in cc file
TVM_FFI_STATIC_INIT_BLOCK({ IntImmNode::RegisterReflection(); });

The RegisterReflection gives us a reflection API to register each member of the object. We can use this function to visit the node and serialize any language object recursively. It also allows us to get members of an object easily in front-end language. For example, we can access the value field of the IntImmNode.

import tvm

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

在不改变前端运行时的情况下，可以在 C++ 中添加新的 Object，这使得对编译器堆栈的扩展变得容易。值得注意的是，这并不是将成员暴露给前端语言的最快捷方式，但可能是最简单的方法之一。还发现这种方法适合我们的目的，因为我们主要使用 Python 进行测试和原型设计，而仍然使用 C++ 来完成繁重的工作。

实现细节#

PackedFunc 中的每个参数都包含联合值 TVMValue 和 type code。这种设计允许动态类型语言直接转换为对应的类型，而静态类型语言则可以在转换过程中进行运行时类型检查。

相关文件

• packed_func.h 用于 C++ API

• c_runtime_api.cc 用于 C API 以及如何提供回调函数。

为了支持扩展类型，采用了注册系统来登记(register)类型的相关信息，比如在 C++ 中的支持情况，详情请见 Extension types 部分。