【TVM源码学习笔记】3 模型编译

 在我们的模型编译运行脚本中，使用relay.build编译模型：

# 设置优化级别
with tvm.transform.PassContext(opt_level=3):
    #编译模型
    lib = relay.build(mod, target, params=params)

因为在python/tvm/relay/__init__.py中有：

from .build_module import build, create_executor, optimize

所以这里build直接调用的是python/tvm/relay/build_module.py中的build函数。我们传入了三个参数，其中第一个mod是一个IRModule实例。excutor参数使用的是默认的graph。函数开头部分的autotvm是TVM的优化工具。所以函数开头部分我们可以先不深究，知道流程走的下面的代码即可：

def build(
    ir_mod,
    target=None,
    target_host=None,
    executor=Executor("graph"),
    runtime=Runtime("cpp"),
    workspace_memory_pools=None,
    constant_memory_pools=None,
    params=None,
    mod_name="default",
):
    ...
 
    with tophub_context:
        bld_mod = BuildModule()
        graph_json, runtime_mod, params = bld_mod.build(
            mod=ir_mod,
            target=raw_targets,
            params=params,
            executor=executor,
            runtime=runtime,
            workspace_memory_pools=workspace_memory_pools,
            constant_memory_pools=constant_memory_pools,
            mod_name=mod_name,
        )
        func_metadata = bld_mod.get_function_metadata()
        devices = bld_mod.get_devices()
        lowered_ir_mods = bld_mod.get_irmodule()
        executor_codegen_metadata = bld_mod.get_executor_codegen_metadata()
 
        ...
 
        elif executor.name == "graph":
            executor_factory = _executor_factory.GraphExecutorFactoryModule(
                ir_mod,
                raw_targets,
                executor,
                graph_json,
                runtime_mod,
                mod_name,
                params,
                func_metadata,
            )
        else:
            assert False, "Executor " + executor + " not supported"
 
        return executor_factory

代码中先实例化了一个BuildModule对象，然后调用BuildModule的build方法编译模型；编译完后读取编译后的模型数据，调用_executor_factory.GraphExecutorFactoryModule创建了一个执行器。

BuildModule.build方法里面，调用的BuildModule._build，这个 _build是在BuildModule.__init__方法中赋值的：

    def __init__(self):
        self.mod = _build_module._BuildModule()
        self._get_graph_json = self.mod["get_graph_json"]
        self._get_module = self.mod["get_module"]
        self._build = self.mod["build"]

self._build挂载的是_build_module.py中_BuildModule()对象的成员。_build_module.py:

import tvm._ffi
 
tvm._ffi._init_api("relay.build_module", __name__)

我们直接搜索relay.build_module._BuildModule，可以找到接口的注册：

runtime::Module RelayBuildCreate() {
  auto exec = make_object();
  return runtime::Module(exec);
}
 
TVM_REGISTER_GLOBAL("relay.build_module._BuildModule").set_body([](TVMArgs args, TVMRetValue* rv) {
  *rv = RelayBuildCreate();
});

这里只是创建了一个RelayBuildModule，然后包再runtime::Module里面，返回runtime::Module实例。

python前端的BuildModule.mod["build"]获取了runtime::Module的build属性，在C++端会首先调用到runtime::Module::GetFunction方法。这个调用过程涉及到TVM的PackedFunc机制，可以参考： 

TVM PackedFunc实现机制 | Don't Respond

runtime::Module::GetFunction的实现：

inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
  return (*this)->GetFunction(name, query_imports);
}

而 runtime::Module对取成员操作符->做了重载:

inline ModuleNode* Module::operator->() { return static_cast(get_mutable()); }

get_mutable是Module祖先类ObjectRef的方法，获取自己的_data字段，是对应的Object实例。而在RelayBuildCreate里面，创建Module实例的时候，传入的是一个RelayBuildModule，继承自Object。所以这里的_data就是这个RelayBuildModule，调用到的(*this)->GetFunction调用到的也就是RelayBuildModule::GetFunction：

class RelayBuildModule : public runtime::ModuleNode {
 public:
 
  ...
 
  PackedFunc GetFunction(const std::string& name, const ObjectPtr