【TVM系列三】算子转换调用流程

文章同步更新在公众号 AIPlayer，欢迎扫码关注，共同进步

 目录

一、前言

二、TVM代码结构

三、 前端流程

四、算子relay转换

五、Relay到TOPI实现的调用过程

六、总结

---

一、前言

本文将从源码分析一个算子在TVM中是如何从前端到后端的转换过程。首先来看一下keras模型编译推理的一个示例。

通过pip安装好keras和tensorflow，按照TVM官方文档的示例进行一些修改，因为环境配置的不同，python包的导入以及编译target进行了修改，示例使用cuda，这里使用cpu。跑通的代码如下：

import tvm
from tvm import te
import tvm.relay as relay
from tvm.contrib.download import download_testdata
from tensorflow import keras
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from tensorflow.keras.applications.resnet50 import preprocess_input
# 下载并导入resnet50模型参数
if tuple(keras.__version__.split(".")) < ("2", "4", "0"):
    weights_url = "".join(
        [
            "https://github.com/fchollet/deep-learning-models/releases/",
            "download/v0.2/resnet50_weights_tf_dim_ordering_tf_kernels.h5",
        ]
    )
    weights_file = "resnet50_keras_old.h5"
else:
    weights_url = "".join(
        [
            " https://storage.googleapis.com/tensorflow/keras-applications/",
            "resnet/resnet50_weights_tf_dim_ordering_tf_kernels.h5",
        ]
    )
    weights_file = "resnet50_keras_new.h5"

weights_path = download_testdata(weights_url, weights_file, module="keras")
keras_resnet50 = keras.applications.resnet50.ResNet50(
    include_top=True, weights=None, input_shape=(224, 224, 3), classes=1000
)                                       
keras_resnet50.load_weights(weights_path)

# 下载一张测试图像并做一些预处理，主要是数据转为float32和layout从NHWC转为NCHW
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
plt.imshow(img)
plt.show()
# input preprocess
data = np.array(img)[np.newaxis, :].astype("float32")
data = preprocess_input(data).transpose([0, 3, 1, 2])   #NHWC -> NCHW   
print("input_1", data.shape) 
shape_dict = {"input_1": data.shape}

# 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()

# 使用编译后的model进行推理得到结果
dtype = "float32"
tvm_out = model(tvm.nd.array(data.astype(dtype)))
tvm_out = tvm_out.numpy()[0]
top1_tvm = np.argmax(tvm_out)

synset_url = "".join(
    [
        "https://gist.githubusercontent.com/zhreshold/",
        "4d0b62f3d01426887599d4f7ede23ee5/raw/",
        "596b27d23537e5a1b5751d2b0481ef172f58b539/",
        "imagenet1000_clsid_to_human.txt",
    ]
)
synset_name = "imagenet1000_clsid_to_human.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
    synset = eval(f.read())
print("Relay top-1 id: {}, prob: {}, class name: {}".format(top1_tvm, tvm_out[top1_tvm], synset[top1_tvm]))
# confirm correctness with keras output
keras_out = keras_resnet50.predict(data.transpose([0, 2, 3, 1]))
keras_out = keras_out[0]
top1_keras = np.argmax(keras_out)
print("Keras top-1 id: {}, prob: {}, class name: {}".format(top1_keras, keras_out[top1_keras], synset[top1_keras]))

在jupyter notebook下运行的结果如下，可以看到，输出的结果一致：

二、TVM代码结构

src/relay 的代码主要处理神经网络的计算图，图节点的编译和执行由src的其它代码实现。python 目录提供了C++ API的封装，图节点对应的算子在src/relay/op中注册，算子的具体实现在topi，通过C++或者Python实现。

当用户使用relay.build(...)执行图编译时，TVM会对图中的每个节点执行下面的动作：

• 在算子注册表中查找算子的实现

• 生成计算表达式和调度

• 将算子编译成目标对象代码

三、 前端流程

以keras前端转换conv2d算子为例，python部分的关键代码调用关系如下图所示：

在keras的示例中，从前端加载模型的调用为：

mod, params = relay.frontend.from_keras(keras_resnet50, shape_dict)

from_keras()函数会调用_convert_layer()，此函数会调用keras_op_to_layers()将keras前端定义的模型layers转换成TVM的Relay表达式，主要是通过一个全局的转换字典_convert_map={layer名称：转换函数}，比如其中的卷积层：{"Conv2D": _convert_convolution}，而_convert_convolution()函数会将卷积相关的参数weights、kernel size、padding以及layout等传入_op.nn.conv2d()函数，而_op.nn.conv2d()会调用_make.conv2d()来运行算子创建函数，主要是通过tvm._ffi._init_api("relay.op.nn._make", __name__)函数实现。那么此函数是如何找到对应的conv2d算子的创建函数呢？。

tvm._ffi._init_api("relay.op.nn._make", __name__)函数调用_init_api_prefix("relay.op.nn._make", ...)，主要是通过sys.modules["relay.op.nn._make"]查找到对应的模块，然后通过list_global_func_names()获取已经注册的全局函数列表，这里调用了C++实现的函数TVMFuncListGlobalNames()，在加载FFI模块的时候，会调用_load_lib()通过ctypes的库函数加载tvm_runtime.so。这个表的内容是一些已经注册的函数名称，在_init_api_prefix()函数中加打印的部分输出如下：

遍历这个表，查找节点对应的函数名，通过_get_global_func(name,...)获取到PackedFunc的控制句柄，调用的也是C++部分的接口_LIB.TVMFuncGetGlobal()。

四、算子relay转换

在上面的Python代码处理过程中，调用了两个C++的接口函数，一个是TVMFuncListGlobalNames()，它主要是返回了Registry注册表中已经注册的函数名称列表，另一个TVMFuncGetGlobal()则是根据name从注册表中查找对应的PackedFunc对象。那么算子是在哪里注册的呢？

在TVM中，Python与C++之间是通过自己实现的PackedFunc来进行连接的，比如conv2d算子，主要是通过三个宏注册相应的接口：

• TVM_REGISTER_NODE_TYPE(Conv2DAttrs)：创建conv2d的attrs对象，在TVM中，算子的计算所需要的参数都以Attribute的方式来定义。

• TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d).set_body_typed()：将"relay.op.nn._make.conv2d注册到注册表中，并调用MakeConv()，MakeConv会返回以op和attrs为参数初始化的Call对象，Call对象的初始化过程中会创建CallNode对象并赋值。

• RELAY_REGISTER_OR("nn.conv2d")：将con2d的OpNode注册到relay的节点注册表中并设置节点的属性参数。

在TVM中一个算子就是通过上面三个宏进行注册的，但是这里只是注册了算子并且设置了算子计算所需要的属性参数，并不涉及算子的计算。从第二小节可以知道同，算子计算实现在topi里。

五、Relay到TOPI实现的调用过程

Relay到TOPI算子的连接是在tvm/python/tvm/relay/op/目录下实现的，根据不同的算子类型归类到几个子目录下，比如vision是处理视觉的算子，image处理图像，nn处理神经网络等等。一般情况下，Relay OP与TOPI OP是通过@reg.register_compute()进行连接，如下：​​​​​​​

@reg.register_compute("nn.upsampling")
def compute_upsampling(attrs, inputs, out_dtype):
    ...

而对于conv2d这种具有多种类型的算子，tvm会使用策略模式的方式来实现Relay与TOPI的连接，调用过程为：

首先通过reg.register_strategy("nn.conv2d", strategy.conv2d_strategy)注册conv2d的处理策略，在conv2d_strategy()中会根据不同的layout或者group参数选择不同的卷积计算策略，通过wrap_compute_conv2d(topi.nn.conv2d_nchw)与topi模块连接，在conv2d_nchw()中会调用te.compute()进行计算，此函数的第二个参数fcompute就是定义的计算规则，它会转为body传到C++的函数TensorComputeOp()或者ComputeOp()进行运算。

六、总结

本文主要从代码层面介绍了TVM算子从keras前端到Relay，再到TOPI算子的转换过程。

往期推荐

【TVM系列一】开发环境搭建

【TVM系列二】TVM介绍