TVM源码阅读一：python与C++相互调用

艺术之光

一、ctypes简介

tvm中有很多地方是python调用C++函数的实现，也有些地方是C++调用python的实现。按照我的理解是有python保证易用性，便于用户使用，C++负责保证性能。
他们之间相互调用归根结底使用的是ctypes的机制，所以这个地方先简单看一下ctypes。
ctypes 是 Python 的外部函数库。它提供了与 C 兼容的数据类型，并允许调用 DLL 或共享库中的函数。可使用该模块以纯 Python 形式对这些库进行封装。
1.1 Python调用C++的简单示例

首先在C++中定义一个全局函数，并编译生成C++动态库。
//mul_module.c
#include <stdio.h>

extern "C"{
    int mul(int a, int b) {
        int c = a * b;
        return c;
    }
}
编译生成动态库，动态库不同的系统后缀不同（Windows 的 dll，Linux 的 so，Mac 的 dylib)
gcc -fPIC -shared mul_module.c -o mul_module.so用ctypes模块在Python中加载生成的动态库（mul_module.so），调用C++中的函数。
import ctypes

# Load shared library
_LIB = ctypes.CDLL("./mul_module.so", ctypes.RTLD_GLOBAL)

a = ctypes.c_int(1)
b = ctypes.c_int(2)
# Call C func in Python
print(_LIB.add(a, b))1.2 C++调用Python的简单示例

在Python中定义函数
def mul(a, b):
  return a * b在C++中为Python的mul定义一个函数原型 int(int, int)
extern "C" {
    typedef int (*PyCFunc)(int, int);
    int call_py_func(PyCFunc f, int a, int b);
}

#include "test.h"
int call_py_func(PyCFunc f, int a, int b) {
  return f(a, b);
}
使用ctypes将Python函数转换成C function，传入C++中进行调用。
import ctypes

#set return and params
cfunc = ctypes.CFUNCTYPE(
    ctypes.c_int, # return type
    ctypes.c_int, # arg0 type
    ctypes.c_int  # arg1 type
    )
#Convert to C function
f = cfunc(add)


# Call Python func in C
print(_LIB.call_py_func(f, 5, 1))
二、PackedFunc

从上面可以看到Python调用C++非常简单，但是C++调用Python是比较麻烦的，先得在C++中定义原型和调用接口，然后进行调用。
TVM呢是想做成一个统一的调用结构，将不同类型的函数包装成统一的函数原型，把这种差异屏蔽掉。
void(TVMArgs args, TVMRetValue *rv);
2.1 几个相关数据结构

TVMValue：存储Python和C++交互的时候支持的类型
typedef union {
  int64_t v_int64;
  double v_float64;
  void* v_handle;
  const char* v_str;
  DLDataType v_type;
  DLDevice v_device;
} TVMValue;
TVMArgs：封装传给PackedFunc的参数
主要包括：TVMValue、参数类型编码、参数个数
class TVMArgs {
public:
  const TVMValue* values;
  const int* type_codes;
  int num_args;
  TVMArgs(const TVMValue* values, const int* type_codes, int num_args)
      : values(values), type_codes(type_codes), num_args(num_args) {}
  /*! \return size of the arguments */
  inline int size() const;
  //重载[]运算符，方便对于多个参数的情况可以通过下标索引直接获取对应的入参
  inline TVMArgValue operator[](int i) const;
};
TVMPODValue_：处理POD类型的数据，由于ctypes在Python和C++之间传递类型、指针、数组有所不同，往往需要类型转换，因此使用TVMPODValue进行了强制类型转换运算符的重载。
需要特别注意的是PackedFunc的强转：会在参数中带上一个所调用函数的指针，此处可参考Ctypes的C++调用Python时的声明中的f，此处就是给函数调用加上f
operator PackedFunc() const {
    if (type_code_ == kTVMNullptr) {
      return PackedFunc(ObjectPtr<Object>(nullptr));
    }
    TVM_CHECK_TYPE_CODE(type_code_, kTVMPackedFuncHandle);
    return PackedFunc(ObjectPtr<Object>(static_cast<Object*>(value_.v_handle)));
  }
TVMArgValue：这个类又继承了TVMPODValue_类，在他的基础上扩充了更多的类型转换的重载，包括string和TypedPackedFunc，传递参数值时使用
operator TypedPackedFunc<FType>() const {
    return TypedPackedFunc<FType>(operator PackedFunc());
  }
TypedPackedFunc的调用会被通过调用PackedFunc转换为带着函数指针的调用
TVMRetValue：这个类也继承了TVMPODValue_类，在他的基础上主要对赋值运算符进行了重载，作为存放调用PackedFunc返回值的容器。
与TVMArgValue二者的区别是TVMRetValue在析构时会释放源数据               
2.2 PackedeFunc

class PackedFunc : public ObjectRef {
public:
  PackedFunc(std::nullptr_t null) : ObjectRef(nullptr) {}  // NOLINT(*)
  //定义参数类型TCallable（TCallable可以强制转换为函数指针std::function<void(TVMArgs, TVMRetValue*)>同时要求TCallable不能是PackedFunc的基类）
  template <typename TCallable,
            typename = std::enable_if_t<
                std::is_convertible<TCallable, std::function<void(TVMArgs, TVMRetValue*)>>::value &&
                !std::is_base_of<TCallable, PackedFunc>::value>>  
  explicit PackedFunc(TCallable data) {
    using ObjType = PackedFuncSubObj<TCallable>;
    //完成对于data的赋值
    data_ = make_object<ObjType>(std::forward<TCallable>(data));
  }
  //通过重载"()"操作符，使得PackedFunc类的对象通过"()"操作，模仿函数调用效果。
  //typename... 表示可变模板
  template <typename... Args>
  inline TVMRetValue operator()(Args&&... args) const;
  TVM_ALWAYS_INLINE void CallPacked(TVMArgs args, TVMRetValue* rv) const;
  /*! \return Whether the packed function is nullptr */
  bool operator==(std::nullptr_t null) const { return data_ == nullptr; }
  /*! \return Whether the packed function is not nullptr */
  bool operator!=(std::nullptr_t null) const { return data_ != nullptr; }

  TVM_DEFINE_OBJECT_REF_METHODS(PackedFunc, ObjectRef, PackedFuncObj);
};
其中主要是通过重载运算符（），在重载函数中打包好参数，将不同输入参数转化为统一的类型无关调用格式std::function<void(TVMArgs, TVMRetValue*)
重载函数的实现如下：
template <typename... Args>
inline TVMRetValue PackedFunc::operator()(Args&&... args) const {
  const int kNumArgs = sizeof...(Args);//sizeof...(Args)表示获取可变参数数量。
  const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
  TVMValue values[kArraySize];
  int type_codes[kArraySize];
  //展开可变参数并使用TVMArgsSetter赋值
  //TVMArgsSetter函数的作用是将调用PackedFunc传入的参数转化为TVMValue类型。
  detail::for_each(TVMArgsSetter(values, type_codes), std::forward<Args>(args)...);
  TVMRetValue rv;
  //获取指针并转换为PackedFuncObj对象
  //接着构造TVMArgs(values, type_codes, kNumArgs)类。将不同输入参数转化为统一的类型无关调用格式
  //传递给CallPacked完成PackedFunc调用。
  (static_cast<PackedFuncObj*>(data_.get()))
      ->CallPacked(TVMArgs(values, type_codes, kNumArgs), &rv);
  return rv;
}
三、函数注册

此处会给一个Python调用C++的例子，主要的一个函数注册宏是TVM_REGISTER_GLOBAL，先简单看一下TVM_REGISTER_GLOBAL和Registry类
3.1 TVM_REGISTER_GLOBAL

#define TVM_REGISTER_GLOBAL(OpName) \
    TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)
一直往下追可以扩展为：
//xxx为递增的counter
::tvm::runtime::Registry& __mk_xxx
    =  ::tvm::runtime::Registry::Register(OpName)
3.2 Registry类

class Registry {
public:
  //三种注册函数的接口
  Registry& set_body(PackedFunc f);
  Registry& set_body_typed(FLambda f);
  Registry& set_body_method(R (T::*f)(Args...));
​  
  //使用name完成调用注册函数实现函数的注册
  static Registry& Register(const std::string& name);
  //在哈希表中寻找名字为name的函数并返回
  static const PackedFunc* Get(const std::string& name);
  //创建函数名列表
  static std::vector ListNames();
​
protected:
  std::string name_;
  PackedFunc func_;
  friend struct Manager;
};
此处额外需要了解的是Manager结构体，主要负责管理注册函数
由三个部分组成：

• 一个map存储函数名和计数
  
• 锁，保证多线程的正确性
  
• 静态函数Global，获取单例inst，限制类的实例化对象个数
  

struct Registry::Manager {
  std::unordered_map<std::string, Registry*> fmap;
  // mutex
  std::mutex mutex;

  Manager() {}

  static Manager* Global() {
    static Manager* inst = new Manager();
    return inst;
  }
};
3.3 自顶向下看Python调用C++函数

以Relay的build为例：
with autotvm.apply_history_best(log_filename):
    with tvm.transform.PassContext(opt_level=3):
        # 编译模型至目标平台，保存在lib变量中，后面可以被导出。
        lib_total = relay.build(sym, target=target, params=params)进入到python/tvm/relay/build_module.py
调用其中的build函数
其中涉及到调用
from . import _build_module
...
self.mod = _build_module._BuildModule()进入到_build_module.py中有
import tvm._ffi

tvm._ffi._init_api("relay.build_module", __name__)把c++ relay.build_module中注册的函数以python PackedFunc对象的形式关联到了_build_module这个模块
再去看_init_api这个函数，这个函数是把注册函数关联到各个模块的关键，python/tvm/_ffi/registry.py：
def _init_api(namespace, target_module_name=None):
    """Initialize api for a given module name

    namespace : str
       The namespace of the source registry

    target_module_name : str
       The target module name if different from namespace
    """
    target_module_name = target_module_name if target_module_name else namespace
    if namespace.startswith("tvm."):
        _init_api_prefix(target_module_name, namespace[4:])
    else:
        _init_api_prefix(target_module_name, namespace)

def _init_api_prefix(module_name, prefix):
    #sys.modules是一个全局字典，每当程序员导入新的模块，sys.modules将自动记录该模块。 当第二次再导入该模块时，python会直接到字典中查找，从而加快了程序运行的速度。
    module = sys.modules[module_name]

    for name in list_global_func_names():
        if not name.startswith(prefix):
            continue

        fname = name[len(prefix) + 1 :]
        target_module = module

        if fname.find(".") != -1:
            continue
        #get_global_func等同于_get_global_func这个函数
        f = get_global_func(name)
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = "TVM PackedFunc %s. " % fname
        #把前面代码构造的python端PackedFunc对象作为属性设置到相应的模块上
        setattr(target_module, ff.__name__, ff)重点在于get_global_func，一直会调用_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)_get_global_func这个函数返回一个python端的PackedFunc对象
首先初始化一个python端的PackedFunc对象
然后使用_LIB.TVMFuncGetGlobal函数，在其中先让_LIB调用ctypes.CDLL获得共享库中的函数，然后在TVMFuncGetGlobal中使用Registry的get获取已经在C++注册好的的函数
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();
}

它的handle成员存储了c++中new出来的PackedFunc对象（以注册函数作为构造参数）的地址，python端的PackedFunc对象的__call__函数调用了c++的TVMFuncCall这个API，handle作为这个API的参数之一，c++端再把handle转成c++的PackedFunc对象来执行，这样就完成了从python端PackedFunc对象的执行到c++端PackedFunc对象的执行的映射。

最后使用make_packed_func返回一个Python端的PackedFunc对象给ff
3.4 C++调用Python

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) #将Python的PackedFunc转化为C++的packfunc
        #先调用C++的注册机制进行注册
        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](args) 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();
}
总结：

其实说一千道一万，TVM的互调机制可以简述为：在C++和Python两边使用了一个统一的函数原型：void(TVMArgs args, TVMRetValue *rv)，这就是PackedFunc机制，实现主要是重载了函数调用运算符“（）”还是非常精巧的，真正的函数体是通过set_body去设置的。
相互调用其实是每次去全局注册函数表中寻找相应的函数名，然后做两种语言之间PackedFunc对象的转换，再去执行。如果是Python调C++那就是将Python端的PackedFunc对象，按照所调用的函数名转成C++端的PackedFunc，然后执行。
接下来会找时间写Relay那部分，了解不多，有什么纰漏欢迎大家批评指正。
参考资料：

月踏：深入理解TVM：Python/C++互调（上）
月踏：深入理解TVM：Python/C++互调（中）
月踏：深入理解TVM：Python/C++互调（下）
https://hjchen2.github.io/2020/01/10/TVM-PackedFunc%E5%AE%9E%E7%8E%B0%E6%9C%BA%E5%88%B6/
【tvm解析】PACKFUNC机制 - 青铜时代的猪 - 博客园
HELLO七仔：TVM的Packed Func实现