TVM中Python/C++接口调用关系

2021-11-06 06:31:26 阅读：175 来源： 互联网

标签：__ PackedFunc const Python c++ TVM python C++ type

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

标签：__,PackedFunc,const,Python,c++,TVM,python,C++,type
来源： https://www.cnblogs.com/wujianming-110117/p/15515897.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

TVM中Python/C++接口调用关系

二．最底层的c++数据结构支撑（围绕c++端PackedFunc）

2. 先看几个注册实例

1. 基本原理

2. 加载TVM动态库

4. python怎么关联c++的PackedFunc

5. 把注册的函数关联到python各个模块

5. 举一个例子