本文将从一个最简单的图像分类模型，了解基本的推理过程，并以此来了解一下 tvm 编译器的基本工作流程，以及对应核心模块的职责
1. 模型的定义
由于tvm支持onnx模型格式，所以我们从 resnet50-v2-7.onnx 开始
模型的输出：https://s3.amazonaws.com/onnx-model-zoo/synset.txt，其实就是1000个图像分类
当我们执行模型推理时，输入一个图像，输出是一个 shape=(1,1000) 的数组，表示这个图片对应1000个分类的概率，其中概率最大的那个分类，就是预测的结果
2. 从一个最基本的推理开始
推理的过程就是把模型，编译成目标设备上可运行的代码，根据输入数据，返回预测结果
如下，说下几个关键的地方：
- 模型转换：relay.frontend.from_onnx()，把模型从 onnx 格式转换成 relay IRModule 格式
- 编译：relay.build()，完成从IRModule到目标设备代码的编译优化，内部包含了 pass 优化、schedule优化
- 推理：输入 module.set_input()，运行 module.run()，获取结果 module.get_output()
其中最核心的就是2这几个步骤，基本包括了tvm最核心的逻辑
# https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx
onnx_model = onnx.load('./resnet50-v2-7.onnx')
img_data = get_test_img()
input_name = "data"
shape_dict = {input_name: img_data.shape}
mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target="llvm", params=params)
module = graph_executor.GraphModule(lib["default"](tvm.device("llvm", 0)))
module.set_input(input_name, tvm.nd.array(img_data.astype('float32')))
module.run()
"https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
out = module.get_output(0, tvm.nd.empty((1, 1000))).numpy()
我们再来看下 relay.build 具体干了啥：
- python/tvm/relay/build_module.py: build()
- python/tvm/relay/build_module.py:BuildModule::build()
- src/relay/backend/build_module.cc: BuildRelay()
- python/tvm/relay/build_module.py:BuildModule::build()
BuildRelay 干的事情就是：Compile a Relay IR module to runtime module
void BuildRelay(IRModule relay_module, const String& mod_name) {
// Relay IRModule -> IRModule optimizations.
IRModule module = WithAttrs(
relay_module, {{tvm::attr::kExecutor, executor_}, {tvm::attr::kRuntime, runtime_}});
relay_module = OptimizeImpl(std::move(module));
// ...
// Generate code for the updated function.
executor_codegen_ = MakeExecutorCodegen(executor_->name);
executor_codegen_->Init(nullptr, config_->primitive_targets);
executor_codegen_->Codegen(func_module, func, mod_name);
//...
auto lowered_funcs = executor_codegen_->GetIRModule();
// When there is no lowered_funcs due to reasons such as optimization.
if (lowered_funcs.size() == 0) {
//...
} else {
ret_.mod = tvm::TIRToRuntime(lowered_funcs, host_target);
}
}
可以看到它包含了 Optimize，Codegen，TIRToRuntime 3个过程
- OptimizeImpl：负责执行一系列的 Pass 优化，什么算符融合，死代码删除，常量折叠，统统都在这里，tvm构建了一个完备和庞大的 Pass infra，可以参考：https://tvm.hyper.ai/docs/arch/arch/pass_infra
- Codegen：生成TIR，这个是最接近目标设备的中间IR了，其中应该还包括一些内存布局相关的优化
- TIRToRuntime：生成目标设备代码，前面不管是 relay IR 还是 TIR，都是以 IRModule 的形式存在的，到这里最终就转成 runtime Module 了，芯片层的对接大部分工作应该在这里
注意：
1）TIRToRuntime 那里重构了好几个版本，以前叫 runtime::Module build()，其实是一样的
2）我在实际调试的过程中，发现FoldConstant这个Pass优化也会调用TIRToRuntime（因为常量折叠优化本身就是要把一些计算前置到编译阶段）
3. 编译的逻辑架构
再来看下官方的逻辑架构图，理解这个图，需要特别注意虚线和实线
- 实线：接口层面的依赖，比如frontends 会依赖 relay.IRModule
- 虚线：表示实现层面的依赖，比如 relay 的pass优化依赖 topi
再回到我们上面整个 build 过程：
- OptimizeImpl：主要包括 relay，topi，te
- Codegen：主要包括 topi，tir
- TIRToRuntime：主要就是 runtime 了
4. TVM Stack
社区有一个图 TVM Stack，画的很好，我们先看下，帮助我们更好的理解 TVM 内部架构
上面是TVM的5层Stack，每一层的核心职责都不一样，下面是TVM内部架构流，描述了TVM核心组件之间的关系
- 第一层：前端 Frameworks，在tvm之前，不同深度学习框架，都有自己的计算图定义，所以 tvm 第一个要干的事情，就是要对接不同的深度学习框架，把他们的图转换成tvm定义的计算图。frontends -> relay
- 第二层：Computation Graph Optimization，图优化，简单来 说就是对图做各种精简。tvm内有一个非常庞大的 pass infra，有各种各样的 pass，pass的输入是图，输出也是一个图，pass不改变图的推理结果（pass是编译器领域一个基本的概念，表示一种优化）
- 第三层：Tensor Compute Description，主要是把图从 relay IR 转变成了更有 numpy 风格的 topi，
- 第四层：Schedule Space and Optimizations，调度优化，主要就是在下面的 tir 这一层
- 第五层：Backends，对接后端各种异构硬件，也就是最底层的IR在不同硬件上的具体实现了