【TVM系列四】模型编译与运行过程

一、前言

针对神经网络模型的编译，TVM封装了非常简洁的python接口，如下：

# keras前端导入，使用llvm作为target编译
mod, params = relay.frontend.from_keras(keras_resnet50, shape_dict)
# compile the model
target = "llvm"
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=0):
    model = relay.build_module.create_executor("graph", mod, dev, target, params).evaluate()
print(model)
tvm_out = model(tvm.nd.array(data.astype(dtype)))

在上一篇文章中介绍了模型的算子转换与Relay IR Module的流程，当TVM将Relay IR Module模型编译为runtime module时，可以通过下面的函数完成：

model = relay.build_module.create_executor("graph", mod, dev, target, params).evaluate()

它返回的是一个函数入口，从打印的输出可以看到它所指向的函数：

<function graphexecutor._make_executor.._graph_wrapper at 0x7fb634b3ed30>

下面这一句相当于将输入传入这个函数去运行：

tvm_out = model(tvm.nd.array(data.astype(dtype)))

这两个步骤在TVM中的运行过程是怎么样的呢？这篇文章将围绕这个问题进行相关的介绍。

二、模型编译

image.png

首先来看一下TVM是如何调用到编译函数的：

• create_executor(...)函数会根据executor的类型返回相应的执行器对象，这里使用的是"graph"，所以返回的是GraphExecutor(...)对象，类class GraphExecutor()是_interpreter.Executor的子类。

• 类class Executor(object)是一个接口类，它的成员函数_make_executor()是一个接口函数，继承它的子类需要实现，它的另一个成员函数evaluate()会调用_make_executor()，也即调用子类class GraphExecutor实现的_make_executor()成员函数。

• _make_executor()主要的工作是调用build(...)函数对Relay IR Module进行编译，并且提供graph执行器的运行函数，也就是前言小节中提到的：

  <function graphexecutor._make_executor.._graph_wrapper at 0x7fb634b3ed30>
  

下面介绍一下build(...)函数的过程：

• 首先实例化一个build_module对象bld_mod，对象的实例化流程是通过tvm._ffi._init_api("relay.build_module", name)调用C++端的接口RelayBuildCreate()，它会创建RelayBuildModule对象。

• 然后通过bld_mod["build"]查找到编译函数的PackedFunc，这个主要是通过类class RelayBuildModule中的GetFunction(...)实现。它会调用类class RelayBuildModule的成员函数Build(...)，在对一些成员变量进行赋值后，调用最终的BuilRelay(...)将Relay IR Module编译为runtime module。

• 最后会返回一个GraphExecutorFactoryModule(...)对象，这个对象在初始化的时候会调用C++端的接口创建类class GraphExecutorFactory的对象。

其中BuildRelay的调用过程如下图所示，它主要有两个步骤，一个是对relay_module做了OptimizeImpl(...)的优化，对模型进行算子的融合、fold constants以及其它的一些优化；第二个是创建codegen对象并生成lowered_funcs并调用tvm::build(...)进行编译。

image.png

整个模型编译的过程可以总结为：将Relay IR Module编译为runtime module并为其构造好执行器的对象。

三、模型运行

模型运行时会调用_make_executor(...)里定义的_graph_wrapper(...)函数：

def _graph_wrapper(*args, **kwargs):
    args = self._convert_args(self.mod["main"], args, kwargs)
# Create map of inputs.
for i, arg in enumerate(args):
        gmodule.set_input(i, arg)
# Run the module, and fetch the output.
    gmodule.run()
    flattened = []
for i in range(gmodule.get_num_outputs()):
        flattened.append(gmodule.get_output(i).copyto(_nd.cpu(0)))
    unflattened = _unflatten(iter(flattened), ret_type)
return unflattened

• 该函数首先会遍历输入的数据，并调用set_input(...)为module设置输入参数，然后调用run()进行模型推理，最后获取输出的结果。

• 在类class GraphExecutorFactory的成员函数GetFunction(...)中，如果输入的name是模型名称，则通过ExecutorCreate(...)创建GraphExecutor对象。

• 在GraphExecutor对象中的GetFunction(...)会根据名称"set_input"与"run"返回相应的PackedFunc对象。

四、总结

本文主要介绍了TVM模型编译与运行过程中的代码流程。