【TVM源码学习笔记】附录1 TVM python调用C++的机制

1. tvm relay op python调用C++

在tvm前端涉及relay算子（比如说外部框架算子转vm relay ir）的时候，会调用到tvm/relay/op下对应算子的接口，进而调用_make.xxx()接口。这个接口最终会调用到C++端对应的算子处理接口。这里我们将探寻Python调用C++的实现。

我们以2D卷积算子为例，对应的接口在python/tvm/relay/op/nn/nn.py中：

def conv2d(
    ...
):
    ...
    return _make.conv2d(
        ...
    )

_make的导入：

from . import _make 

也就是python/tvm/relay/op/nn/_make.py：

import tvm._ffi

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

这里__name__是一个python内置变量，表示当前模块的文件名（不包括.py），即tvm/relay/op/nn/_make。

tvm._ffi模块位于python/tvm/_ffi。函数_init_api的定义在python/tvm/_ffi/registry.py中

def _init_api(namespace, target_module_name=None):
    """Initialize api for a given module name

    namespace : str
       The namespace of the source registry

    target_module_name : str
       The target module name if different from namespace
    """
    target_module_name = target_module_name if target_module_name else namespace
    if namespace.startswith("tvm."):
        _init_api_prefix(target_module_name, namespace[4:])
    else:
        _init_api_prefix(target_module_name, namespace)

这里传入的第一个参数namespace为relay.op.nn._make， target_module_name参数为tvm/relay/op/nn/_make。这样传入_init_api_prefix的两个参数将是 tvm.relay.op.nn._make和relay.op.nn._make。

def _init_api_prefix(module_name, prefix):
    module = sys.modules[module_name]

    for name in list_global_func_names():
        if not name.startswith(prefix):
            continue

        fname = name[len(prefix) + 1 :]
        target_module = module

        if fname.find(".") != -1:
            continue
        f = get_global_func(name)
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = "TVM PackedFunc %s. " % fname
        setattr(target_module, ff.__name__, ff)

module = sys.modules[module_name]获取的是tvm.relay.op.nn._make模块的句柄。list_global_func_names()定义在python/tvm/_ffi/registry.py中：

def list_global_func_names():
    """Get list of global functions registered.

    Returns
    -------
    names : list
       List of global functions names.
    """
    plist = ctypes.POINTER(ctypes.c_char_p)()
    size = ctypes.c_uint()

    check_call(_LIB.TVMFuncListGlobalNames(ctypes.byref(size), ctypes.byref(plist)))
    fnames = []
    for i in range(size.value):
        fnames.append(py_str(plist[i]))
    return fnames

接口种通过ctypes方式，调用C++库的TVMFuncListGlobalNames接口，得到的结果字符串数组plist，该数组为所有全局接口的函数名集合。TVMFuncListGlobalNames接口定义在src/runtime/registry.cc中

int TVMFuncListGlobalNames(int* out_size, const char*** out_array) {
  API_BEGIN();
  TVMFuncThreadLocalEntry* ret = TVMFuncThreadLocalStore::Get();
  ret->ret_vec_str = tvm::runtime::Registry::ListNames();
  ret->ret_vec_charp.clear();
  for (size_t i = 0; i < ret->ret_vec_str.size(); ++i) {
    ret->ret_vec_charp.push_back(ret->ret_vec_str[i].c_str());
  }
  *out_array = dmlc::BeginPtr(ret->ret_vec_charp);
  *out_size = static_cast(ret->ret_vec_str.size());
  API_END();
}

函数中调用tvm::runtime::Registry::ListNames()得到函数名表：

std::vector Registry::ListNames() {
  Manager* m = Manager::Global();
  std::lock_guard lock(m->mutex);
  std::vector keys;
  keys.reserve(m->fmap.size());
  for (const auto& kv : m->fmap) {
    keys.push_back(kv.first);
  }
  return keys;
}

可以看到，函数名都是从Manager类实例的fmap表的第一个元素。而且Manager还是各单实例类。而fmap的定义：

struct Registry::Manager {
  // map storing the functions.
  // We deliberately used raw pointer.
  // This is because PackedFunc can contain callbacks into the host language (Python) and the
  // resource can become invalid because of indeterministic order of destruction and forking.
  // The resources will only be recycled during program exit.
  std::unordered_map fmap;
  // mutex
  std::mutex mutex;

  Manager() {}

  static Manager* Global() {
    // We deliberately leak the Manager instance, to avoid leak sanitizers
    // complaining about the entries in Manager::fmap being leaked at program
    // exit.
    static Manager* inst = new Manager();
    return inst;
  }
};

从注释看，这个fmap是一个存储函数的map表。表单元的第一个元素是string类型。

再看_init_api_prefix中的get_global_func：

def get_global_func(name, allow_missing=False):
    return _get_global_func(name, allow_missing)

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)

_get_global_func中使用ctypes方式调用C++库中的TVMFuncGetGlobal函数：

int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
  API_BEGIN();
  const tvm::runtime::PackedFunc* fp = tvm::runtime::Registry::Get(name);
  if (fp != nullptr) {
    *out = new tvm::runtime::PackedFunc(*fp);  // NOLINT(*)
  } else {
    *out = nullptr;
  }
  API_END();
}

const PackedFunc* Registry::Get(const std::string& name) {
  Manager* m = Manager::Global();
  std::lock_guard lock(m->mutex);
  auto it = m->fmap.find(name);
  if (it == m->fmap.end()) return nullptr;
  return &(it->second->func_);
}

TVMFuncGetGlobal调用了Registry::Get，从Manager的fmap表中，找到第一个元素为python传入的函数名的单元，从该单元的第二个元素中获取了函数指针。也就是根据函数名获取函数句柄。

 搜索下谁在往fmap成员中写数据，可以看到是Registry::Register接口：

Registry& Registry::Register(const std::string& name, bool can_override) {  // NOLINT(*)
  Manager* m = Manager::Global();
  std::lock_guard lock(m->mutex);
  if (m->fmap.count(name)) {
    ICHECK(can_override) << "Global PackedFunc " << name << " is already registered";
  }

  Registry* r = new Registry();
  r->name_ = name;
  m->fmap[name] = r;
  return *r;
}

可以看到调用Registry::Register接口接口时，如果name在fmap中不存在，就会创建一个Registry实例，加入Manager的fmap表，并返回新建的Registry实例。搜索Registry::Register接口的调用，在include/tvm/runtime/registry.h中有定义

/*!
 * \brief Register a function globally.
 * \code
 *   TVM_REGISTER_GLOBAL("MyPrint")
 *   .set_body([](TVMArgs args, TVMRetValue* rv) {
 *   });
 * \endcode
 */
#define TVM_REGISTER_GLOBAL(OpName) \
  TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)

这里调用Registry::Register接口，传入的是一个函数名。在代码中搜索TVM_REGISTER_GLOBAL宏的使用会有很多。这里我们继续关注relay.op.nn._make.conv2d的，搜索到src/relay/op/nn/convolution.cc中代码：

TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d")
    .set_body_typed([](Expr data, Expr weight, Array strides, Array padding,
                       Array dilation, int groups, IndexExpr channels,
                       Array kernel_size, String data_layout, String kernel_layout,
                       String out_layout, DataType out_dtype) {
      return MakeConv(data, weight, strides, padding, dilation, groups, channels,
                                   kernel_size, data_layout, kernel_layout, out_layout, out_dtype,
                                   "nn.conv2d");
    });

这里set_body_typed的参数为一个lamad表达式,函数体部分调用了MakeConv。所以这里是向Manager的fmap注册了一个函数，名字为relay.op.nn._make.conv2d， 函数体部分是调用MakeConv。

简单的讲，TVM relay op的python到C++调用，就是在C++里，创建一个函数管理表（Manager::fmap），各算子向这个表注册接口，并给每个接口一个标记符；python部分在对应目录下放一个_make.py文件，在这个文件中设置注册的函数标记符和对应的C++函数句柄；当算子转换以标记符调用接口时，就会调用到C++里面的对应的函数体。

我们看下MakeConv的实现：

template 
inline Expr MakeConv(Expr data, Expr weight, Array strides, Array padding,
                     Array dilation, int groups, IndexExpr channels,
                     Array kernel_size, std::string data_layout,
                     std::string kernel_layout, std::string out_layout, DataType out_dtype,
                     std::string op_name) {
  auto attrs = make_object();
  attrs->strides = std::move(strides);
  attrs->padding = std::move(padding);
  attrs->dilation = std::move(dilation);
  attrs->groups = groups;
  attrs->channels = std::move(channels);
  attrs->kernel_size = std::move(kernel_size);
  attrs->data_layout = std::move(data_layout);
  attrs->kernel_layout = std::move(kernel_layout);
  attrs->out_layout = std::move(out_layout);
  attrs->out_dtype = std::move(out_dtype);
  const Op& op = Op::Get(op_name);
  return Call(op, {data, weight}, Attrs(attrs), {});
}

 这里将卷积参数打包，生成一个Op实例，然后生成一个Call实例返回。

python/tvm/relay/op/tensor.py中的数学算符可能会多绕一步，根据算符的类型，先在src/relay/op/tensor目录下的文件中TVM_REGISTER_NODE_TYPE、RELAY_REGISTER_BINARY_OP或者RELAY_REGISTER_UNARY_OP定义，这些宏中再调用TVM_REGISTER_GLOBAL宏，参见src/relay/op/op_common.h

2. tvm relay function python调用C++

在Graphproto.from_onnx的最后，使用网络的输入输出和权重参数打包成一个Function实例，然后生成一个IRModule实例：

# 由模型输入, 输出表达式依赖的权重和输出表达式生成一个function
func = _function.Function([v for k, v in self._inputs.items()], outputs)
# 返回表达式和所有权重
return IRModule.from_expr(func), self._params

这两步也都是会调用到C++代码。先看_function.Function的流程：

@tvm._ffi.register_object("relay.Function")
class Function(BaseFunc):
    def __init__(self, params, body, ret_type=None, type_params=None, attrs=None):
        if type_params is None:
            type_params = convert([])

        self.__init_handle_by_constructor__(
            _ffi_api.Function, params, body, ret_type, type_params, attrs
        )

    def __call__(self, *args):
        return Call(self, args, None, None)

__init__函数第二个参数body是函数体，而前面在调用_function.Function时传入的时outputs。这是因为outputs并不是网络或者函数的输出张量，而是输出的计算表达式，而且这个表达式描述的是从输入开始，一步一步的到输出的计算过程，也就是函数实现的所有计算过程了。所以这个outputs就是函数体。

__init__中调用了self.__init_handle_by_constructor__，参数_ffi_api.Function这种形式在前面算子调用流程中我们已经分析过，_ffi_api引入的是模块，Function是具体的函数，所以我们看下当前目录下的_ffi_api是什么模块， 见python/tvm/relay/_ffi_api.py：

import tvm._ffi

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

模块为relay.ir，所以_ffi_api.Function就是relay.ir.Function。

搜索该标记符的注册TVM_REGISTER_GLOBAL("relay.ir.Function")可以看到：

TVM_REGISTER_GLOBAL("relay.ir.Function")
    .set_body_typed([](tvm::Array params, Expr body, Type ret_type,
                       tvm::Array ty_params, tvm::DictAttrs attrs) {
      return Function(params, body, ret_type, ty_params, attrs);
    });

也就是调用_ffi_api.Function会在C++端实例化一个Function。

在python的Function类中， _ffi_api.Function是作为参数传给self.__init_handle_by_constructor__，这个方法定义在python/tvm/_ffi/_ctypes/object.py中的基类ObjectBase中，而ObjectBase.__init_handle_by_constructor__调用的是

def __init_handle_by_constructor__(fconstructor, args):
    """Initialize handle by constructor"""
    temp_args = []
    values, tcodes, num_args = _make_tvm_args(args, temp_args)
    ret_val = TVMValue()
    ret_tcode = ctypes.c_int()
    if (
        _LIB.TVMFuncCall(
            fconstructor.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
    assert ret_tcode.value == ArgTypeCode.OBJECT_HANDLE
    handle = ret_val.v_handle
    return handle

我们看下TVMFuncCall的调用链

src/runtime/c_runtime_api.cc:

int TVMFuncCall(TVMFunctionHandle func, TVMValue* args, int* arg_type_codes, int num_args,
                TVMValue* ret_val, int* ret_type_code) {
  API_BEGIN();

  TVMRetValue rv;
  (*static_cast(func)).CallPacked(TVMArgs(args, arg_type_codes, num_args), &rv);

...

}


include/tvm/runtime/packed_func.h:

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

这里最后调用到的body_就是TVM_REGISTER_GLOBAL("relay.ir.Function").set_typed_body设置的lamabd函数体。

这里比较绕，我们理下：

1. 首先将注册的relay.ir.Function作为参数传给了__init_handle_by_constructor__；

2. __init_handle_by_constructor__调用了_LIB.TVMFuncCall；

3. _LIB.TVMFuncCall相当于一个函数执行器，它执行了relay.ir.Function；

4. relay.ir.Function的函数体被执行时，返回一个C++端的Function对象句柄。

3. tvm relay op IRModule python调用C++

onnx.py中GraphProto.from_onnx最后return IRModule.from_expr(func), self._params，这个from_expr代码在python/tvm/ir/module.py中：

    def from_expr(expr, functions=None, type_defs=None):
        funcs = functions if functions is not None else {}
        defs = type_defs if type_defs is not None else {}
        return _ffi_api.Module_FromExpr(expr, funcs, defs)

这里直接调用_ffi_api.Module_FromExpr，python/tvm/ir/目录定义的模块名为ir(见python/tvm/ir/_ffi_api.py)， 搜索对应的函数注册TVM_REGISTER_GLOBAL("ir.Module_FromExpr")，注册函数执行IRModule::FromExpr，FromExpr调用IRModule::FromExprInContext，生成一个C++端的IRModule实例