【tvm解析】PACKFUNC机制

为实现多种语言支持，需要满足以下几点：

• 部署：编译结果可以从python/javascript/c++调用。
• Debug: 在python中定义一个函数，在编译函数中调用。
• 链接：编写驱动程序以调用设备特定代码（如CUDA），可以在编译的host侧调用
• ‎原型：python侧定义IR PASS，并从C++后端调用该代码‎
• 接口暴露：c++后端代码暴露到python侧
• ‎实验：将编译的函数运送到嵌入式设备，可以直接在嵌入式设备上运行

tvm希望在任何一个语言中定义的函数，可以在其他的语言中都可以调用。同样希望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);
}

上面的例子中，定义了一个MyAdd的PackedFunc，接受两个参数，args表示输入参数， rv表示返回值。这个参数是类型无关的(type-erased)，这意味着函数签名中对输入输出参数的类型没有限制。这样，当调用这个函数的时候， 从栈上获取输入参数（TVMArgs），通过TVMRetValue返回函数返回值。

通过C++的模板技巧，可以像正常函数一样调用PackedFunc。由于类型无关的特性，可以在像python这样的动态类型的语言中调用PackedFunc，而无需插入额外其他的胶水代码。下面展示了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

这个限制，使得实现及其简单而且无需序列化操作。虽然增加了限制，但对于DL开发来说，大多数场景下仅仅需要传递DLTensor和数字就够了。

既然PackedFunc可以将另外的PackedFunc作为函数参数，那就可以在python与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。

PackFunction不仅用于tvm编译器中，同样也用于开发的技术栈中。在tvm中所有的PASS函数都通过PackedFunc暴露给前端的。编译结果同样是通过PackedFunc打包的。

为了保证runtime尽可能的小，runtime中隔离了IR对象的支持。这使得runtime大小只有200~600k，具体的大小取决于平台驱动部分。

PackedFunc带来的调用开销很小，仅仅是通过栈传递了一些参数对象，只要不通过它包装较小的函数，就是OK的。总之，PackedFunc是tvm中通用的胶水代码，支持了tvm的编译部署。

额外的部分：

c++ 注册，python调用

上文中介绍注册时，使用到了一个C++宏TVM_REGISTER_GLOBAL，这里介绍中间是如何链接起来的。

TVM_REGISTER_GLOBAL("callhello")
.set_body([](TVMArgs args, TVMRetValue* rv) {
  PackedFunc f = args[0];
  f("hello world");
});

//展开就是
TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register("callhello").set_body([](TVMArgs args, TVMRetValue* rv) {
  PackedFunc f = args[0];
  f("hello world");
});

这里的::tvm::runtime::Registry::Register

Registry& Registry::Register(const std::string& name, bool can_override) {  // NOLINT(*)
  Manager* m = Manager::Global();//这是个静态对象，Manager持有一个map来记录注册对象
  std::lock_guard<std::mutex> lock(m->mutex);
  if (m->fmap.count(name)) {
    ICHECK(can_override) << "Global PackedFunc " << name << " is already registered";
  }

  Registry* r = new Registry();
  r->name_ = name;
  m->fmap[name] = r;
  return *r;
}

下面看下Registry的实现。

/*! \brief Registry for global function */
class Registry {
 public:
  //设置函数体
  TVM_DLL Registry& set_body(PackedFunc f);  // NOLINT(*)
  Registry& set_body(PackedFunc::FType f) {  // NOLINT(*)
    return set_body(PackedFunc(f));
  }

  //给一个任意函数，萃取函数签名
  template <typename FLambda>
  Registry& set_body_typed(FLambda f) {
    using FType = typename detail::function_signature<FLambda>::FType;
    return set_body(TypedPackedFunc<FType>(std::move(f), name_).packed());
  }
  //给一个类成员函数、返回值、参数，使用lambda包装
  template <typename T, typename R, typename... Args>
  Registry& set_body_method(R (T::*f)(Args...)) {
    auto fwrap = [f](T target, Args... params) -> R {
      // call method pointer
      return (target.*f)(params...);
    };
    return set_body(TypedPackedFunc<R(T, Args...)>(fwrap, name_));
  }

  template <typename T, typename R, typename... Args>
  Registry& set_body_method(R (T::*f)(Args...) const) {
    auto fwrap = [f](const T target, Args... params) -> R {
      // call method pointer
      return (target.*f)(params...);
    };
    return set_body(TypedPackedFunc<R(const T, Args...)>(fwrap, name_));
  }
  //
  template <typename TObjectRef, typename TNode, typename R, typename... Args,
            typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
  Registry& set_body_method(R (TNode::*f)(Args...)) {
    auto fwrap = [f](TObjectRef ref, Args... params) {
      TNode* target = ref.operator->();
      // call method pointer
      return (target->*f)(params...);
    };
    return set_body(TypedPackedFunc<R(TObjectRef, Args...)>(fwrap, name_));
  }

  template <typename TObjectRef, typename TNode, typename R, typename... Args,
            typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
  Registry& set_body_method(R (TNode::*f)(Args...) const) {
    auto fwrap = [f](TObjectRef ref, Args... params) {
      const TNode* target = ref.operator->();
      // call method pointer
      return (target->*f)(params...);
    };
    return set_body(TypedPackedFunc<R(TObjectRef, Args...)>(fwrap, name_));
  }

  TVM_DLL static Registry& Register(const std::string& name, bool override = false);  // NOLINT(*)

  TVM_DLL static bool Remove(const std::string& name);

  TVM_DLL static const PackedFunc* Get(const std::string& name);
  TVM_DLL static std::vector<std::string> ListNames();

  struct Manager;

 protected:
  std::string name_;
  PackedFunc func_;
  friend struct Manager;
};

上面注册以后是在一个全局对象中，下一部就看python侧如何调用的。

python端最终会调用到 _get_global_func函数，具体实现如下。

def _get_global_func(name, allow_missing=False):
    handle = PackedFuncHandle()
    check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))

    if handle.value:
        return _make_packed_func(handle, False)

    if allow_missing:
        return None

    raise ValueError("Cannot find global function %s" % name)

进而会调用到TVMFuncGetGlobal

int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
  API_BEGIN();
  const tvm::runtime::PackedFunc* fp = tvm::runtime::Registry::Get(name);
  if (fp != nullptr) {
    *out = new tvm::runtime::PackedFunc(*fp);  // NOLINT(*)
  } else {
    *out = nullptr;
  }
  API_END();
}

这里既可以发现tvm::runtime::Registry::Get(name)来查找相关注册函数的。

python注册，c++ 调用

如下面的函数，通过装饰器注册。

@tvm._ffi.register_func("relay.backend.lower_call")

在c++中调用

static auto flower_call = tvm::runtime::Registry::Get("relay.backend.lower_call");

下面介绍以下python的注册。

def register_func(func_name, f=None, override=False):
    if callable(func_name):
        f = func_name
        func_name = f.__name__

    if not isinstance(func_name, str):
        raise ValueError("expect string function name")

    ioverride = ctypes.c_int(override)

    def register(myf):
        """internal register function"""
        if not isinstance(myf, PackedFuncBase):
            myf = convert_to_tvm_func(myf) #转化为packfunc
        #注册
        check_call(_LIB.TVMFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride))
        return myf

    if f:
        return register(f)
    return register

def convert_to_tvm_func(pyfunc):
    local_pyfunc = pyfunc

    def cfun(args, type_codes, num_args, ret, _):
        """ ctypes function """
        num_args = num_args.value if isinstance(num_args, ctypes.c_int) else num_args
        pyargs = (C_TO_PY_ARG_SWITCH[type_codes[i]](args[i]) for i in range(num_args))
        # pylint: disable=broad-except
        try:
            rv = local_pyfunc(*pyargs)
        except Exception:
            msg = traceback.format_exc()
            msg = py2cerror(msg)
            _LIB.TVMAPISetLastError(c_str(msg))
            return -1

        if rv is not None:
            if isinstance(rv, tuple):
                raise ValueError("PackedFunction can only support one return value")
            temp_args = []
            values, tcodes, _ = _make_tvm_args((rv,), temp_args)
            if not isinstance(ret, TVMRetValueHandle):
                ret = TVMRetValueHandle(ret)
            if _LIB.TVMCFuncSetReturn(ret, values, tcodes, ctypes.c_int(1)) != 0:
                raise get_last_ffi_error()
            _ = temp_args
            _ = rv
        return 0

    handle = PackedFuncHandle()
    f = TVMPackedCFunc(cfun)
    # NOTE: We will need to use python-api to increase ref count of the f
    # TVM_FREE_PYOBJ will be called after it is no longer needed.
    pyobj = ctypes.py_object(f)
    ctypes.pythonapi.Py_IncRef(pyobj)
    if _LIB.TVMFuncCreateFromCFunc(f, pyobj, TVM_FREE_PYOBJ, ctypes.byref(handle)) != 0:
        raise get_last_ffi_error()
    return _make_packed_func(handle, False)

int TVMFuncRegisterGlobal(const char* name, TVMFunctionHandle f, int override) {
  API_BEGIN();
  tvm::runtime::Registry::Register(name, override != 0)
      .set_body(*static_cast<tvm::runtime::PackedFunc*>(f));
  API_END();
}