TVM中PythonC++接口调用关系

TVM中Python/C++接口调用关系
TVM已经是一个很庞大的系统，包含了很多的功能模块，python和c++的互相调用这个功能模块，没有使用第三方的开源库（boost.python、pybind11等），自主实现了一套复杂但精致高效强大的机制。这部分内容很多，本文分成三部分，分析实现原理：
最底层的c++数据结构支撑（围绕c++端PackedFunc）
基于PackedFunc的函数注册（围绕TVM_REGISTER_GLOBAL）
偏上层的python的调用细节（围绕ctypes内置库和python端PackedFunc）
一．最底层的c++数据结构支撑（围绕c++端PackedFunc）

1. 概述
   PackedFunc类是python和c++调用关系的桥梁，此类实现代码在include/tvm/runtime/packed_func.h文件中，这里有一个TypedPackedFunc类，只是PackedFunc的一个wrapper，主要增加了类型检查的功能，开发TVM的c++代码，要尽可能的使用这个类，但是为了把问题尽可能的简化，只关注PackedFunc这个最底层类，其中用到了下面这几个关键的数据结构：
   • TVMValue
   • TVMArgs
   • TVMPODValue_
   • TVMArgValue
   • TVMRetValue
   • TVMArgsSetter
   2.TVMValue
   这是最基本的一个数据结构，是一个union，主要是为了储存c++和其它语言交互时所支持的几种类型的数据，代码很简单（其中DLDataType和DLDevice是两个复合数据类型，可以到github查看细节）：
   // include/tvm/runtime/c_runtime_api.h
   typedef union {
   int64_t v_int64;
   double v_float64;
   void* v_handle;
   const char* v_str;
   DLDataType v_type;
   DLDevice v_device;
   } TVMValue;
   3.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) { … }

inline int size() const { return num_args; }
inline TVMArgValue operator[](int i) const {
return TVMArgValue(values[i], type_codes[i]);
}
};
4.TVMPODValue_
这是一个内部使用的基类，主要主要服务于后面介绍到的TVMArgValue和TVMRetValue，这个类主要是处理POD类型的数据，POD是plain old data的缩写，是scalar type，trival type，standard layout type三者之一，具体可参考cppreference的PODType、is_pod、is_scalar、is_trivial、is_standard_layout等。这个类的实现核心是强制类型转换运算符重载（在c++中，类型的名字，包括类的名字本身也是一种运算符，即类型强制转换运算符），如下面代码所示：
class TVMPODValue_ {
public:
operator double() const { return value_.v_float64; }
operator int64_t() const { return value_.v_int64; }
operator void*() const { return value_.v_handle; }
template
T* ptr() const { return static_cast(value_.v_handle); }

protected:
TVMValue value_;
int type_code_;
};
5.TVMArgValue
这个类继承TVMPODValue_类，用作表示PackedFunc的一个参数，与TVMPODValue_的区别是扩充了一些数据类型的支持，如string、PackedFunc、TypedPackedFunc等，对后两个的支持是在c++代码中能够调用python函数的根本原因。这个类只使用所保存的underlying data，不会去做释放，代码如下：
class TVMArgValue : public TVMPODValue_ {
public:
TVMArgValue() {}
TVMArgValue(TVMValue value, int type_code)
: TVMPODValue_(value, type_code) {}

operator std::string() const {}
operator PackedFunc() const { return *ptr(); }
const TVMValue& value() const { return value_; }

template
inline operator T() const;
inline operator DLDataType() const;
inline operator DataType() const;
};
6.TVMRetValue
这个类也是继承自TVMPODValue_类，主要作用是作为存放调用PackedFunc返回值的容器，与TVMArgValue的区别是，会管理所保存的underlying data，会做释放。这个类主要由四部分构成：
• 构造和析构函数
• 对强制类型转换运算符重载的扩展
• 对赋值运算符的重载
• 辅助函数，包括释放资源的Clear函数
代码如下：
class TVMRetValue : public TVMPODValue_ {
public:
// ctor and dtor, dtor will release related buffer
TVMRetValue() {}
~TVMRetValue() { this->Clear(); }

// conversion operators
operator std::string() const { return *ptrstd::string(); }
operator DLDataType() const { return value_.v_type; }
operator PackedFunc() const { return *ptr(); }

// Assign operators
TVMRetValue& operator=(double value) {}
TVMRetValue& operator=(void* value) {}
TVMRetValue& operator=(int64_t value) {}
TVMRetValue& operator=(std::string value) {}
TVMRetValue& operator=(PackedFunc f) {}

private:
// judge type_code_, release underlying data
void Clear() {
if (type_code_ == kTVMStr || type_code_ == kTVMBytes) {
delete ptrstd::string();
} else if(type_code_ == kTVMPackedFuncHandle) {
delete ptr();
} else if(type_code_ == kTVMNDArrayHandle) {
NDArray::FFIDecRef(
static_cast(value_.v_handle));
} else if(type_code_ == kTVMModuleHandle
|| type_code_ == kTVMObjectHandle ) {
static_cast(value_.v_handle)->DecRef();
}
type_code_ = kTVMNullptr;
}
};
7.TVMArgsSetter
这是一个用于给TVMValue对象赋值的辅助类，主要通过重载函数调用运算符来实现，主要实现原理如下：
class TVMArgsSetter {
public:
TVMArgsSetter(TVMValue* values, int* type_codes)
: values_(values), type_codes_(type_codes) {}

void operator()(size_t i, double value) const {
values_[i].v_float64 = value;
type_codes_[i] = kDLFloat;
}
void operator()(size_t i, const string& value) const {
values_[i].v_str = value.c_str();
type_codes_[i] = kTVMStr;
}
void operator()(size_t i, const PackedFunc& value) const {
values_[i].v_handle = const_cast(&value);
type_codes_[i] = kTVMPackedFuncHandle;
}
private:
TVMValue* values_;
int* type_codes_;
};
8.PackedFunc
有了前面所述的数据结构作为基础，再来看PackedFunc的实现，PackedFunc的实现很简单，内部只使用了一个储存函数指针的变量，再通过重载函数调用运算符，调用这个函数指针所指向的函数，代码如下：
class PackedFunc {
public:
using FType = function;
PackedFunc() {}
explicit PackedFunc(FType body) : body_(body) {}

template
inline TVMRetValue operator()(Args&&… args) const {
const int kNumArgs = sizeof…(Args);
const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
TVMValue values[kArraySize];
int type_codes[kArraySize];
detail::for_each(TVMArgsSetter(values, type_codes),
std::forward(args)…);
TVMRetValue rv;
body_(TVMArgs(values, type_codes, kNumArgs), &rv);
return rv;
}

inline void CallPacked(TVMArgs args, TVMRetValue* rv) const {
body_(args, rv);
}

private:
FType body_;
};
9.小结
TVM的官方文档对PackedFunc机制有一段简短精辟的介绍（https://tvm.apache.org/docs/dev/runtime.html），大家可以作为参考来理解上面代码：
PackedFunc 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 TVMArgs on stack, and gets the result back via TVMRetValue. Thanks to template tricks in C++, we can call a PackedFunc just like a normal function. Because of its type-erased nature, we can call a PackedFunc from dynamic languages like python, without additional glue code for each new type function created.
二．最底层的c++数据结构支撑（围绕c++端PackedFunc）

1. 概述
   前面已经讲过python/c++调用关系的c++端的底层核心数据结构：PackedFunc。本节是python/c++调用关系这个系列的第二部分，主要来讲c++端的函数注册，python端对c++端的函数调用，都来源于c++端的注册函数，最主要的一个函数注册宏是TVM_REGISTER_GLOBAL，code base里，大概用了1300多次，除了这个注册宏，TVM里还有许多其它的注册宏。
2. 先看几个注册实例
   注册的函数可以是普通函数，也可以是labda表达式，注册接口有三个：set_body、set_body_typed、set_body_method，第一个使用的是PackedFunc，后面两个使用的是TypedPackedFunc，PackedFunc在这个系列的前面讲过了，TypedPackedFunc是PackedFunc的一个wrapper，实现比较复杂，暂时不介绍。下面举三个简单示例，展示下这三个注册接口的使用。
   使用set_body接口注册lambda表达式：
   // src/topi/nn.cc
   TVM_REGISTER_GLOBAL(“topi.nn.relu”)
   .set_body([](TVMArgs args, TVMRetValue* rv) {
   *rv = relu(args[0]);
   });
   使用set_body_typed接口注册lambda表达式：
   // src/te/schedule/graph.cc
   TVM_REGISTER_GLOBAL(“schedule.PostDFSOrder”)
   .set_body_typed([](
   const Array& roots,
   const ReadGraph& g) {
   return PostDFSOrder(roots, g);
   });
   使用set_body_method接口注册类内函数：
   // src/ir/module.cc
   TVM_REGISTER_GLOBAL(“ir.Module_GetGlobalVar”)
   .set_body_method(&IRModuleNode::GetGlobalVar);
   //
3. TVM_REGISTER_GLOBAL宏定义
   这个宏定义的本质，就是在注册文件定义了一个static的引用变量，引用到注册机内部new出来的一个新的Registry对象：
   // include/tvm/runtime/registry.h
   #define TVM_REGISTER_GLOBAL(OpName)
   static ::tvm::runtime::Registry& __mk_TVMxxx =
   ::tvm::runtime::Registry::Register(OpName)
   上面的xxx其实是__COUNTER__这个编译器拓展宏，生成的一个唯一标识符，GCC文档里对这个宏有详细的描述。（https://gcc.gnu.org/onlinedocs/cpp/Common-Predefined-Macros.html）：
   This macro expands to sequential integral values starting from 0. In conjunction with the ## operator, this provides a convenient means to generate unique identifiers. Care must be taken to ensure that COUNTER is not expanded prior to inclusion of precompiled headers which use it. Otherwise, the precompiled headers will not be used.
4. Registry::Manager
   先来看最核心的Manager类，是Registry的内部类，用来存储注册的对象，先看下代码：
   // src/runtime/registry.cc
   struct Registry::Manager {
   static Manager* Global() {
   static Manager* inst = new Manager();
   return inst;
   }
   std::mutex mutex;
   unordered_map fmap;
   };
   这个数据结构很简单，从上面代码能得到下面几点信息：
   • 数据结构里面带锁，可以保证线程安全。
   • Manager是个单例，限制类的实例化对象个数是一种技术，可以限制实例化对象个数为0个、1个、N个，具体可参照《More Effective C++：35个改善编程与设计的有效方法》的条款26：限制某个 class 所能产生的对象数量这一章节。
   • 使用unordered_map存储注册信息，注册对象是Registry指针。
5. Registry
   这才是注册机的核心数据结构，简化过的代码如下（只保留了关键的数据结构和接口，原文使用了大量的模板、泛型等c++用法）：
   // include/tvm/runtime/registry.h
   class Registry {
   public:
   Registry& set_body(PackedFunc f);
   Registry& set_body_typed(FLambda f);
   Registry& set_body_method(R (T::*f)(Args…));

static Registry& Register(const std::string& name);
static const PackedFunc* Get(const std::string& name);
static std::vector ListNames();

protected:
std::string name_;
PackedFunc func_;
friend struct Manager;
};
Registry的功能可以为三部分，相关的实现代码也比较简单，总结如下：
• 设置注册函数的set_body系列接口，使用Registry的一系列set_body方法，可以把PackedFunc类型的函数对象，设置到Registry对象中。
• 创建Registry对象的Register静态接口，参照下面代码：
Registry& Registry::Register(const std::string& name) {
Manager* m = Manager::Global();
std::lock_guardstd::mutex lock(m->mutex);

Registry* r = new Registry();
r->name_ = name;
m->fmap[name] = r;
return r;
}
获取注册函数的Get静态接口，代码如下：
const PackedFunc Registry::Get(const std::string& name) {
Manager* m = Manager::Global();
std::lock_guardstd::mutex lock(m->mutex);
auto it = m->fmap.find(name);
if (it == m->fmap.end()) return nullptr;
return &(it->second->func_);
}

6. 小结
   对于python/c++的调用关系至关重要，注册机也是一个所有深度学习框架、编译器都会用到的技术，很有必要了解清楚。
   三．偏上层的python的调用细节（围绕ctypes内置库和python端PackedFunc）
   

7. 基本原理
   TVM使用python的ctypes模块，调用c++代码提供的API，ctypes是python内建的可以用于调用C/C++动态链接库函数的功能模块，ctypes官方文档（https://docs.python.org/3/library/ctypes.html）是这样介绍的：
   ctypes is a foreign function library for Python.It provides C compatible data types, and allows calling functions in DLLs or shared libraries. It can be used to wrap these libraries in pure Python.
   对于动态链接库提供的API，需要使用符合c语言编译和链接约定的API，因为python的ctype只和c兼容，c++编译器会对函数和变量名进行name mangling，使用__cplusplus宏和extern "C"得到符合c语言编译和链接约定的API，以TVM给python提供的接口为例：
   // TVM给python提供的接口主要都在这个文件：
   // include/tvm/runtime/c_runtime_api.h，
   // 下面主要展示了__cplusplus和extern "C"的用法，
   // 以及几个关键的API。
   #ifdef __cplusplus
   extern “C” {
   #endif

int TVMFuncListGlobalNames(…);
int TVMFuncGetGlobal(…);
int TVMFuncCall(…);

#ifdef __cplusplus
} // TVM_EXTERN_C
#endif
2. 加载TVM动态库
TVM的python代码从python/tvm/init.py中开始真正执行，即：
from ._ffi.base import TVMError, version
这句简单的import代码，会执行python/tvm/_ffi/init.py：
from .base import register_error
from .registry import register_func
from .registry import _init_api, get_global_func
上面的第一句，会导致python/tvm/_ffi/base.py中的下面代码被执行：
def _load_lib():
lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL)
return lib, os.path.basename(lib_path[0])

_LIB, _LIB_NAME = _load_lib()
上面的lib_path[0]是TVM动态链接库的全路径名称，在linux系统做的试验，链接库的名称是/xxx/libtvm.so（不同的系统动态库的名字会有所不同，windows系统是.dll，苹果系统是.dylib，linux系统是.so），在_load_lib函数执行完成后，_LIB和_LIB_NAME都完成了初始化，其中_LIB是一个ctypes.CDLL类型的变量，可以认为能够 *** 作TVM动态链接库的export symbols的一个全局句柄，_LIB_NAME是libtvm.so这个字符串。这样后续在python中，就能通过_LIB这个桥梁，不断与c++的部分进行交互。
4. python怎么关联c++的PackedFunc
前面已经对c++中的PackedFunc做了详细的剖析，这里主要理清楚python的代码中，怎么使用这个核心组件的，还是通过代码，一步步来看。
python中获取c++API的底层函数是_get_global_func：

def _get_global_func(func_name):
handle = ctypes.c_void_p()
_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle))
return _make_packed_func(handle, False)
这里面handle是一个相当于void类型的指针变量，因为从ctypes的官方文档中可以查到，c_void_p对应的primitive C compatible data type是：

_get_global_func中调用了TVMFuncGetGlobal这个API，从这个API的实现发现，handle最终保存了一个c++代码在堆中new出来的PackedFunc对象指针：
// src/runtime/registry.cc
int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
const tvm::runtime::PackedFunc* fp
= tvm::runtime::Registry::Get(name);
*out = new tvm::runtime::PackedFunc(*fp);
}
和c++PackedFunc的关联工作这时候才完成一半，在_get_global_func的最后调用了_make_packed_func这个函数：

def _make_packed_func(handle, is_global):
obj = PackedFunc.new(PackedFuncbase)
obj.is_global = is_global
obj.handle = handle
return obj
可以看到_make_packed_func函数中，创建了一个定义在python/tvm/runtime/packed_func.py中的python PackedFunc对象，PackedFunc其实是一个空实现，继承自PackedFuncbase类，PackedFuncbase类中定义了一个__call__函数：

class PackedFuncbase(object):
def call(self, args):
values, tcodes, num_args = _make_tvm_args(args, temp_args)
ret_val = TVMValue()
ret_tcode = ctypes.c_int()
_LIB.TVMFuncCall(
self.handle,
values,
tcodes,
ctypes.c_int(num_args),
ctypes.byref(ret_val),
ctypes.byref(ret_tcode),
)
return ret_val
从上面可以看出，python的__call__函数，调用了C的TVMFuncCall这个API，把前面保存有c++ PackedFunc对象地址的handle，以及相关的函数参数传了进去，TVMFuncCall的主体代码如下：
// src/runtime/c_runtime_api.cc
int TVMFuncCall(TVMFunctionHandle handle, TVMValue args, …)
(static_cast>(handle))
.CallPacked(TVMArgs(args, arg_type_codes, num_args), &rv);
}
这样就完成了把c++中的PackedFunc映射到了python中的PackedFunc，在python代码中只需要调用python中创建好的PackedFunc对象，就会通过上面分析的过程，一步步调到c++的代码中。
5. 把注册的函数关联到python各个模块
注册的函数既包括c++中注册的函数，也包括python中注册的函数，主要是c++中注册的函数，通过list_global_func_names函数（实际上调用的TVMFuncListGlobalNames这个c++API），可以得到c++中注册的所有函数，目前有1500多个，截图了最开始的十个作为示例，显示一下：

for name in list_global_func_names():
    if not name.startswith(prefix):
        continue
    fname = name[len(prefix) + 1 :]
    f = get_global_func(name)
    ff = _get_api(f)
    ff.__name__ = fname
    ff.__doc__ = "TVM PackedFunc %s. " % fname
    setattr(target_module, ff.__name__, ff)