TVM python调用c++函数

TVM使用c++进行算子编译和部署，使用python语言做编译器前端，用来调用c++编写的代码，这样做的好处是既利用了python的易用性又利用了c++的性能，但是对于新手，看tvm代码时会非常的迷惑。

这篇文章就我自己的理解来解释tvm是如何使用这两种语言的，水平有限，欢迎大佬批评指正。

目录结构

首先来看tvm仓库的目录结构，其中最主要的子目录有这么三个：

• include：c++的头文件，定义了算子编译和部署的类和函数
• src：c++的源文件，实现了算子编译和部署的类和函数
• python：python前端，包装了c++的函数和对象

include和src都是c++文件，需要对它们进行编译以构建共享库（在linux系统上是libtvm.so和libtvm_runtime.so文件，在macos上是libtvm.dylib和libtvm_runtime.dylib文件，在windows上是libtvm.dll和libtvm_runtime.dll文件），以下假设在linux系统上，构建的过程在构建共享库（官方文档）。

在python中会对该构建的共享库进行加载，以下我们以一个例子来对python调用c++中定义的函数进行解释。

一个简单的代码示例

首先在c++的src目录中定义一个简单的加法函数，并对其进行注册：

#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>

void MyAdd(tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
  // 自动将参数转换为所需的类型。
  int a = args[0];
  int b = args[1];
  // 自动赋值返回给 rv
  *rv = a + b;
}

// 在 C++ 中注册一个全局打包函数
TVM_REGISTER_GLOBAL("myadd").set_body(MyAdd);

因为我们要在python中对这个函数进行调用，因此我们需要重新构建.so动态库。

构建成功后，我们就可以在python目录下使用python对该函数进行调用了。

import tvm

myadd = tvm.get_global_func("myadd")
# 打印 3
print(myadd(1, 2))

这里我们可能会奇怪，是在哪里对.so动态库进行的调用呢？接下来我们对代码进行分析。

分析代码示例

因为python是前端，因此整个执行过程由python代码进行调用。python的代码非常简单，只有三行，我们按执行顺序分别对其进行解释。

import tvm

在import tvm的时候会首先执行tvm库的初始化文件tvm/__init__.py：

from ._ffi.base import TVMError, __version__, _RUNTIME_ONLY

首先python解释器会在_ffi包下寻找tvm/_ffi/__init__.py文件并执行：

from .base import register_error

因此执行tvm/_ffi/base.py文件，并将register_error导入当前命名空间，在其中最重要的是这么一行：

_LIB, _LIB_NAME = _load_lib()

调用了_load_lib()，它的实现也在该文件中：

def _load_lib():
    """通过搜索可能的路径加载库文件"""
    lib_path = libinfo.find_lib_path() # 搜索库文件的路径
    # The dll search path need to be added explicitly in windows after python 3.8
    if sys.platform.startswith("win32") and sys.version_info >= (3, 8):
        for path in libinfo.get_dll_directories():
            os.add_dll_directory(path)
	# 使用ctypes模块的CDLL类加载c语言编写的动态链接库（DLL）
	# lib_path[0]是参数路径，ctypes.RTLD_GLOBAL表示将加载的符号添加到全局符号表中
    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL)
    lib.TVMGetLastError.restype = ctypes.c_char_p
    return lib, os.path.basename(lib_path[0])

libinfo.find_lib_path()对.so共享库可能在路径进行搜索，这里寻找的路径有环境变量PATH、用户主目录、以及该tvm目录下的一些位置（如build目录）。如果找到了.so共享库，返回找到的路径（一般就是\path\to\libtvm.so和\path\to\libtvm_runtime.so），然后通过ctypes模块的CDLL类加载c语言编写的动态链接库（DLL），返回的CDLL对象lib最终保存为全局变量_LIB。

在加载共享库时执行了共享库的代码，共享库的c++代码是以什么顺序执行的呢？根据我的调试结果，我发现是以源文件名称的字典序执行的，也就是说，先执行src/arith/analyzer.cc中的函数注册，再执行src/arith/bound_deducer.cc中的函数注册，再执行src/arith/const_int_bound.cc中的函数注册……，从而也会执行我们定义的MyAdd函数的注册代码TVM_REGISTER_GLOBAL("myadd").set_body(MyAdd);。

c++端基于PackedFunc的函数注册在这篇文章中解释的很好。

注意_LIB是c语言而不是c++，tvm提供一个最小的c api（位于include/tvm/runtime/c_runtime_api.h），其中使用extern "C"来告诉c++编译器不要对括号内的函数使用c++的名字修饰规则，使用该c api调用关键的c++函数：

#ifdef __cplusplus
extern "C" {
#endif

// 通过函数名调用函数
TVM_DLL int TVMFuncCall(TVMFunctionHandle func, TVMValue* arg_values, int* type_codes, int num_args, TVMValue* ret_val, int* ret_type_code);

// 通过函数名获取某个全局函数
TVM_DLL int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out);

// 列出所有的全局函数名
TVM_DLL int TVMFuncListGlobalNames(int* out_size, const char*** out_array);

#ifdef __cplusplus
} // TVM_EXTERN_C
#endif

其中我们可以看到，tvm通过TVMFunctionHandle来调用一个c++函数，它是一个PackedFunc的句柄（类型为void*），c++端的底层核心数据结构在这篇文章中进行了详细介绍，总体来说，使用PackedFunc的目的是为了类型擦除，也就是函数签名没有限制要传入的输入参数的数量和类型和返回值类型，当我们调用 PackedFunc 时，它将输入参数打包到堆栈上的 TVMArgs，并通过 TVMRetValue 返回结果，从而不需要为新的函数类型添加额外的代码。

因此我们可以看到，import tvm获取了.so共享库并保存到了全局变量_LIB中，以后通过该变量来获取c++中定义的函数。

到此为止，import tvm这句代码做了什么我们已经解释完了，总结来说就是执行了c++共享库，并加载了其中的c语言函数。

myadd = tvm.get_global_func(“myadd”)

接下来解释第二行代码myadd = tvm.get_global_func("myadd")，首先调用到了tvm/_ffi/registry.py文件下的get_global_func()函数，该函数直接调用了tvm/_ffi/_ctypes/packed_func.py中的_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)

其中最重要的是_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle))，这里调用了tvm提供的c api，通过TVMFuncGetGlobal函数得到了想要的函数名name的函数指针handle，然后将该handle包装为PackedFunc并返回。

print(myadd(1, 2))

最后就是python中对myadd函数的调用print(myadd(1, 2))，这段代码调用了PackedFunc函数myadd，对PackedFunc函数的调用是什么过程呢？准确来说在python中PackedFunc是一个类，它位于python/tvm/runtime/packed_func.py，它是一个空实现，继承自PackedFuncBase类（位于python/tvm/_ffi/_ctypes/packed_func.py），该类的实现如下：

class PackedFuncBase(object):

    def __init__(self, handle, is_global):
        """使用句柄初始化函数"""
        self.handle = handle
        self.is_global = is_global

    def __del__(self):...

    def __call__(self, *args):
		"""使用位置参数调用函数"""
        temp_args = []
        values, tcodes, num_args = _make_tvm_args(args, temp_args)
        ret_val = TVMValue()
        ret_tcode = ctypes.c_int()
        if (
            _LIB.TVMFuncCall(
                self.handle,
                values,
                tcodes,
                ctypes.c_int(num_args),
                ctypes.byref(ret_val),
                ctypes.byref(ret_tcode),
            )
            != 0
        ):
            raise get_last_ffi_error()
        _ = temp_args
        _ = args
        return RETURN_SWITCH[ret_tcode.value](ret_val)

myadd(1, 2)调用了PackedFunc类的__call__函数，在该函数中有_LIB.TVMFuncCall(self.handle, values, tcodes, ctypes.c_int(num_args), ctypes.byref(ret_val), ctypes.byref(ret_tcode))，这里调用了tvm提供的c api，通过TVMFuncCall函数调用了c++中实现的MyAdd函数，其中六个参数分别是函数指针、参数、参数类型、参数数量、返回值指针、返回值类型，调用结束后将返回值转换为返回值的类型并返回。

到此为止，python调用c++的代码示例就解释完了。