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 <typename T>

  T* ptr() const { return static_cast<T*>(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<PackedFunc>(); }

  const TVMValue& value() const { return value_; }

  template <typename T>

  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 *ptr<std::string>(); }

  operator DLDataType() const { return value_.v_type; }

  operator PackedFunc() const { return *ptr<PackedFunc>(); }

  // 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 ptr<std::string>();

    } else if(type_code_ == kTVMPackedFuncHandle) {

      delete ptr<PackedFunc>();

    } else if(type_code_ == kTVMNDArrayHandle) {

      NDArray::FFIDecRef(

        static_cast<TVMArrayHandle>(value_.v_handle));

    } else if(type_code_ == kTVMModuleHandle

        || type_code_ == kTVMObjectHandle ) {

      static_cast<Object*>(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<PackedFunc*>(&value);

     type_codes_[i] = kTVMPackedFuncHandle;

  }

 private:

  TVMValue* values_;

  int* type_codes_;

};

8.PackedFunc

有了前面所述的数据结构作为基础，再来看PackedFunc的实现，PackedFunc的实现很简单，内部只使用了一个储存函数指针的变量，再通过重载函数调用运算符，调用这个函数指针所指向的函数，代码如下：

class PackedFunc {

 public:

  using FType = function<void(TVMArgs args, TVMRetValue* rv)>;

  PackedFunc() {}

  explicit PackedFunc(FType body) : body_(body) {}

  template <typename... Args>

  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>(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<float>(args[0]);

});

使用set_body_typed接口注册lambda表达式：

// src/te/schedule/graph.cc

TVM_REGISTER_GLOBAL("schedule.PostDFSOrder")

    .set_body_typed([](

          const Array<Operation>& roots,

          const ReadGraph& g) {

      return PostDFSOrder(roots, g);

    });

使用set_body_method接口注册类内函数：

// src/ir/module.cc

TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar")

    .set_body_method<IRModule>(&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<std::string, Registry*> 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_guard<std::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_guard<std::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）

1. 基本原理

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：

# python/tvm/_ffi/_ctypes/packed_func.py

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这个函数：

# python/tvm/_ffi/_ctypes/packed_func.py

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__函数：

# python/tvm/_ffi/_ctypes/packed_func.py

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<const PackedFunc*>(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多个，截图了最开始的十个作为示例，显示一下：

 先看_init_api这个函数，这个函数是把注册函数关联到各个模块的关键：

# python/tvm/_ffi/registry.py

def _init_api(prefix, module_name):

    target_module = sys.modules[module_name]

​

    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)

这里面有三个最主要的点：

• line3：sys.modules是一个全局字典，每当程序员导入新的模块，sys.modules将自动记录该模块。 当第二次再导入该模块时，python会直接到字典中查找，从而加快了程序运行的速度。
• line9：get_global_func等同于上面介绍的_get_global_func这个函数，这个函数返回一个python端的PackedFunc对象，handle成员存储了c++中new出来的PackedFunc对象（以注册函数作为构造参数）的地址，python端的PackedFunc对象的__call__函数，调用了c++的TVMFuncCall这个API，handle作为这个API的参数之一，c++端再把handle转成c++的PackedFunc对象执行，这样就完成了从python端PackedFunc对象的执行到c++端PackedFunc对象的执行的映射。
• line13：把前面代码构造的python端PackedFunc对象作为属性，设置到相应的模块上。

然后，各个模块中对_init_api全局调用一次，就完成了关联，在代码中找了几个作为示例，如下所示：

# python/tvm/runtime/_ffi_api.py

tvm._ffi._init_api("runtime", __name__)

​

# python/tvm/relay/op/op.py

tvm._ffi._init_api("relay.op", __name__)

​

# python/tvm/relay/backend/_backend.py

tvm._ffi._init_api("relay.backend", __name__)

5. 举一个例子

以TVM中求绝对值的函数abs为例，这个函数实现在tir模块，函数的功能很简单，不会造成额外的理解负担，只关注从python调用，怎么映射到c++中的，先看在c++中abs函数的定义和注册：

// src/tir/op/op.cc

// 函数定义

PrimExpr abs(PrimExpr x, Span span) { ... }

​

// 函数注册

TVM_REGISTER_GLOBAL("tir.abs").set_body_typed(tvm::abs);

再看python端的调用：

# python/tvm/tir/_ffi_api.py

# 把c++ tir中注册的函数以python PackedFunc

# 对象的形式关联到了_ffi_api这个模块

tvm._ffi._init_api("tir", __name__)

​

# python/tvm/tir/op.py

# 定义了abs的python函数，内部调用了前面

# 关联到_ffi_api这个模块的python PackedFunc对象

def abs(x, span=None):

    return _ffi_api.abs(x, span)

最后，用户可以这样使用这个函数：

import tvm

from tvm import tir

​

rlt = tir.abs(-100)

print("abs(-100) = %d" % (rlt)

参考链接：

https://zhuanlan.zhihu.com/p/363991566

https://zhuanlan.zhihu.com/p/365795292

https://www.136.la/jingpin/show-123101.html