2

wkcn · 9月8日

在EAIDK-310上学习TVM (一) - “远程过程调用”例子的详细介绍

在EAIDK-310上学习TVM (一) - “远程过程调用”例子的详细介绍

大家好~在上一篇文章《在EAIDK-310上部署TVM Runtime》里,我介绍了什么是TVM,如何在EAIDK-310开发板上部署TVM Runtime,以及运行了一个简单的远程过程调用例子。
“远程过程调用”例子是TVM官方提供的一个例子,为了将这个例子适配到EAIDK-310开发板,我对例子的代码进行了简单的修改,并将代码上传在Github项目TVM tutorials for EAIDK,大家可以直接下载这个项目。对于这个例子的运行,请参考《在EAIDK-310上部署TVM Runtime》。本篇文章将会详细分析这个例子。

代码可以在cross_compilation_and_rpc.py 链接中查看。

导入模块

# 导入NumPy模块
import numpy as np

# 导入TVM模块,和TVM的远程过程调用模块rpc, 工具模块util. 在本例子中,使用了工具模块util获得临时文件夹的地址。
import tvm
from tvm import rpc
from tvm.contrib import util
# 导入配置模块,用于保存开发板的IP地址与开发板上TVM远程过程调用服务的端口,以及针对开发板的编译参数。官方例子中没有config模块,而是直接填写IP地址和端口等参数。
import config

声明一个核函数

我们需要声明一个核函数,它的作用是输入一个长度为1024的张量(Tensor),返回这个张量逐元素加1.0的结果。

# 创建一个常数节点,值为1024, 保存在变量n中
n = tvm.convert(1024)

# 建立一个张量A,它同时也是一个占位符号,可以作为函数的输入,并且要求输入的张量的形状(shape)为(n, )。因为n是值为1024的常数节点,所以这里要求输入的张量的长度为1024。
A = tvm.placeholder((n,), name='A')

# 建立一个新的张量B,它的形状为(n, ),它是通过lambda函数A[i] + 1.0计算得到的,也就是将张量A逐元素加1.0的结果保存到张量B。
B = tvm.compute((n,), lambda i: A[i] + 1.0, name='B')

# 为张量B的算子建立一个调度,保存到变量s中
s = tvm.create_schedule(B.op)

编译核函数

声明好核函数后,可以对函数进行编译。当我们需要在本地机器执行核函数时,使用默认编译参数即可。当我们需要在远程机器执行核函数,而且远程机器的架构和本地机器不一致时,需要进行交叉编译。

# local_demo表示是否在本机编译,如果local_demo为True则在本机(通常是x86架构)编译;如果local_demo为False则对远程机器(比如ARM架构的EAIDK-310开发板)进行交叉编译(因为编译代码的机器和远程机器不是同一台机器,并且架构不同)。
local_demo = False

# 如果是本地机器,使用llvm默认编译参数编译,否则使用交叉编译参数编译。在config.py文件中可以看到target_host='llvm -target=aarch64-linux-gnu'。因为目标机器(EAIDK-310)的架构是ARMv8架构,运行64位Linux操作系统。
if local_demo:
    target = 'llvm'
else:
    target = config.target_host

# 设置编译签名(signature), 其中s是前面声明的核函数的输出算子的调度。[A, B]指函数的输入和输出,注意这里**没有规定**前面的变量一定是输入,后面的变量一定是输出。变量是输入还是输出取决于核函数的声明,这里也可以写成[B, A],此时就是输出在前输入在后了。target指的是编译参数。name设置了编译出的函数的名字。`tvm.build`函数将会返回一个编译了的函数对象。
func = tvm.build(s, [A, B], target=target, name='add_one')

导出目标文件到本地机器的临时文件夹

将核函数编译后,编译后的函数对象保存在变量func中。我们需要将编译出的目标文件保存到本地机器的临时文件夹中,并且上传到远程机器(开发板)上。

# 获取本机的临时文件夹路径
temp = util.tempdir()
# 我们将目标文件命名为`lib.tar`,获取它保存到临时文件夹下的绝对路径
path = temp.relpath('lib.tar')
# 将目标文件导出的临时文件夹,命名为lib.tar
func.export_library(path)

建立会话

在运行核函数前,我们需要建立会话。会话分为“本地会话”和“远程会话”:“本地会话”用于在本地机器进行调试,“远程会话”用于在远程机器执行核函数。有了会话,同一套代码能很方便地在本地或远程机器上执行。

if local_demo:
    # 对于本地机器,建立本地会话
    remote = rpc.LocalSession()
else:
    # 对于远程机器,建立远程会话,其中host为远程机器的IP地址,port为远程机器的TVM远程过程调用服务使用的端口。
    host = config.host
    port = config.port
    remote = rpc.connect(host, port)

上传目标文件

建立会话后,通过会话上传目标文件lib.tar,并读取目标文件.

# 上传目标文件lib.tar
remote.upload(path)
# 读取目标文件,得到编译后的核函数
func = remote.load_module('lib.tar')

执行核函数

通过会话上传目标文件后,可以读取目标文件并执行核函数。

# 设置上下文为会话中的机器的CPU,上下文指的是运行核函数的环境,比如CPU,GPU.
ctx = remote.cpu()
# 使用NumPy创建一个长度为1024,类型与张量A一致的随机数组,并将这个数组转换为TVM中的多维数组,上下文为ctx.
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
# 建立一个全部为0, 长度为1024的TVM多维数组
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
# 调用核函数 
func(a, b)
# 对比结果
np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)

当我们使用本地机器调试核函数时,数组a和数组b都在本地机器上,直接在本地机器执行核函数得到结果,再通过调用asnumpy函数,在本地机器中将TVM的多维数组转为NumPy的多维数组;当我们使用远程机器,也就是开发板运行核函数时,数组a和数组b将在远程机器上创建,并在远程机器上运行,通过调用asnumpy函数,将远程机器上的多维数组通过网络通信传给本地机器,在本地机器上转为NumPy多维数组。

性能分析

我们可以对核函数的性能进行评估,比如统计函数的运行时间。

# 核函数的time_evaluator将会返回一个评估核函数运行时间的函数,这个时间是实际执行函数的时间,不包括网络通信的时间。测试次数number为10.
time_f = func.time_evaluator(func.entry_name, ctx, number=10)
# 得到实际实行的时间,并计算均值
cost = time_f(a, b).mean
print('%g secs/op' % cost)

总结

在本篇文章中,我分析了TVM的官方例程“远程过程调用”的代码,分为了多个部分:核函数的声明、核函数的编译、导出目标文件、建立会话、上传目标文件、核函数的执行与性能测试。建议大家重点理解“核函数的声名”部分,它声明了整个核函数的计算过程。

大家如果有什么疑问,欢迎在评论区留言:)

2阅读 61
2 条评论
关注数
25
文章数
32
分享EAIDK最新资讯、案例和干货
目录
qrcode
关注微信服务号
实时接收新的回答提醒和评论通知