Python教程

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

本文主要是介绍TVM中Python/C++接口调用关系,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

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

 

这篇关于TVM中Python/C++接口调用关系的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!