TVM：交叉编译和RPC

TVM：交叉编译和RPC

之前我们介绍了 TVM 的安装、本机demo和树莓派远程demo。本文将介绍了在 TVM 中使用 RPC 进行交叉编译和远程设备执行。

通过交叉编译和 RPC，我们可以在本地机器上编译程序，然后在远程设备上运行它。 当远程设备资源有限时很有用，例如 Raspberry Pi 和移动平台。 在本文中，我们将使用 Raspberry Pi 作为 CPU 示例，使用 Firefly-RK3399 作为 OpenCL 示例。

在远程设备上构建 TVM Runtime

首先我们要在远程设备上编译安装 TVM Runtime，注意这里我们对模型的编译是在本机进行的，而远程设备只需要运行模型即可，因此只需要构建 TVM Runtime。

注意：本节和下一节中的所有指令都应在目标设备上执行，例如树莓派。 我们假设它运行着 Linux。

git clone --recursive https://github.com/apache/tvm tvm
cd tvm
make runtime -j2

将 Python 路径添加到环境变量：

export PYTHONPATH=$PYTHONPATH:/path/to/tvm/python

在远程设备上设置 RPC 服务器

在远程设备（如本例中的树莓派）上运行以下命令来开启 RPC 服务器：

python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090

如果看到下面这行说明远程设备上的 RPC 服务已经成功开启了：

INFO:root:RPCServer: bind to 0.0.0.0:9090

在本机上声明并交叉编译核

注意：现在我们回到本地机器了，之后的操作都是在本机（含有完整的，带有 LLVM 的 TVM）上进行。

我们现在本机上声明一个简单的核：

import numpy as np

import tvm
from tvm import te
from tvm import rpc
from tvm.contrib import utils

n = tvm.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
s = te.create_schedule(B.op)

然后我们来对核进行交叉编译。对于树莓派3B来说，target 应该是 llvm -mtriple=armv7l-linux-gnueabihf’ 。如果真的有一个远程设备树莓派的话可以将下面的 local_demo 改为 False ，否则还是保留为 True 使得本 demo 可以正常运行。

local_demo = True

if local_demo:
    target = "llvm"
else:
    target = "llvm -mtriple=armv7l-linux-gnueabihf"

func = tvm.build(s, [A, B], target=target, name="add_one")
# 在本地的临时目录下保存一个 lib
temp = utils.tempdir()
path = temp.relpath("lib.tar")
func.export_library(path)

注意：要使用真正的远程设备运行本教程，请将 local_demo 更改为 False 并将 build 中的 target 替换为适合我们设备的目标三元组（target triple）。不同设备的（target triple）可能不同。例如，对于 Raspberry Pi 3B，它是'llvm -mtriple=armv7l-linux-gnueabihf'，对于 RK3399，它是'llvm -mtriple=aarch64-linux-gnu'。

通常，我们可以通过在您的设备上运行 gcc -v 来查询目标，并查找以 Target: 开头的行（尽管它可能仍然是一个宽松的配置。）

除了 -mtriple，您还可以设置其他编译选项，例如：

• -mcpu=<cpuname>
  指定当前架构中的特定芯片以为其生成代码。默认情况下，这是从 target triple 推断出来的，并自动检测到当前架构。

• -mattr=a1,+a2,-a3,…
  覆盖或控制目标的特定属性，例如是否启用 SIMD 操作。默认属性集由当前 CPU 设置。要获取可用属性列表，我们可以执行以下操作：

llc -mtriple=\<your device target triple\> -mattr=help

这些选项与 llc 一致。建议将 target triple和功能集设置为包含可用的特定功能，以便我们可以充分利用板的功能。可以到 LLVM 交叉编译指南中找到有关交叉编译属性的更多详细信息。

通过RPC在远程运行CPU核

接下来是如何将生成的CPU核运行在远程设备上，首先我们建立与远程设备的 RPC 会话。

if local_demo:
    remote = rpc.LocalSession()
else:
    # 下面的IP是笔者的，请大家换成自己的远程设备的IP
    host = "10.206.105.111"
    port = 9090
    remote = rpc.connect(host, port)

将 lib 上传到远程设备，然后调用设备本地编译器来重新链接它们。 现在 func 是一个远程模块对象。

remote.upload(path)
func = remote.load_module("lib.tar")

# 在远程设备上创建数组
dev = remote.cpu()
a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), dev)
b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
# func 将会运行在远程设备上
func(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)

如果我们想要在远程设备上评估核的性能时，我们需要避免网络开销。 time_evaluator 将返回一个远程函数，该函数多次运行该函数，测量远程设备上每次运行的成本并返回测量的成本。并将网络开销排除在外。

time_f = func.time_evaluator(func.entry_name, dev, number=10)
cost = time_f(a, b).mean
print("%g secs/op" % cost)

此处输出：

1.178e-07 secs/op

通过RPC在远程设备上运行OpenCL核

对于远程 OpenCL 设备，整个流程和上面几乎是一样的。我们定义自己的和核，上传文件，并通过 RPC 运行。

注意：树莓派并不支持 OpenCL，以下代码是在 Firefly-RK3399 上进行测试的。大家可以通过这个教程来为 RK3399 配置操作系统和 OpenCL。

同样我们需要再 RK3399 上构建 TVM Runtiime（注意要在 config.cmake 中启用 OpenCL），在 tvm 根目录下，执行：

cp cmake/config.cmake .
sed -i "s/USE_OPENCL OFF/USE OPENCL ON" config.cmake
make runtime -j4

接下来。我们通过以下代码来远程运行 OpenCL 核：

def run_opencl():
    # 注意，这里是我自己的 RK3399 板子的设置，你可以根据自己的环境进行调整
    opencl_device_host = "10.77.1.145"
    opencl_device_port = 9090
    target = tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")

    # 为上述 'add one' 计算声明创建 schedule
    s = te.create_schedule(B.op)
    xo, xi = s[B].split(B.op.axis[0], factor=32)
    s[B].bind(xo, te.thread_axis("blockIdx.x"))
    s[B].bind(xi, te.thread_axis("threadIdx.x"))
    func = tvm.build(s, [A, B], target=target)

    remote = rpc.connect(opencl_device_host, opencl_device_port)

    # 导出并上传
    path = temp.relpath("lib_cl.tar")
    func.export_library(path)
    remote.upload(path)
    func = remote.load_module("lib_cl.tar")

    # 运行
    dev = remote.cl()
    a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), dev)
    b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
    func(a, b)
    np.testing.assert_equal(b.numpy(), a.numpy() + 1)
    print("OpenCL test passed!")

总结

本文提供了 TVM 中交叉编译和 RPC 功能的演示。

• 在远程设备上设置 RPC 服务器
• 设置设备的 target配置 并在本机上交叉编译核
• 通过 RPC 远程上传并运行核