交叉编译和RPC#

作者: Ziheng Jiang, Lianmin Zheng

本教程介绍了 TVM 中的交叉编译和使用 RPC 的远程设备执行。

通过交叉编译和 RPC,你可以在本地机器上编译程序,然后在远程设备上运行它。在远程设备资源有限的情况下,如 Raspberry Pi 和移动平台,这很有用。在本教程中,将使用 Raspberry Pi 作为 CPU 例子,使用 Firefly-RK3399 作为 OpenCL 例子。

在设备上建立 TVM 运行时#

第一步是在远程设备上建立 TVM 运行时。

注意

本节和下一节的所有指令都应该在目标设备上执行,例如 Raspberry Pi。假设目标设备运行的是 Linux。

由于在本地机器上进行编译,远程设备只用于运行生成的代码。只需要在远程设备上构建 TVM 运行时。

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

成功构建运行时后,需要在 ~/.bashrc 文件中设置环境变量。可以用 vi ~/.bashrc 编辑 ~/.bashrc,并添加以下一行(假设 TVM 目录 ~/tvm)。

export PYTHONPATH=$PYTHONPATH:~/tvm/python

要更新环境变量,执行 source ~/.bashrc

在设备上设置 RPC 服务器#

要启动 RPC 服务器,请在你的远程设备上运行以下命令(本例中为 Raspberry Pi)。

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

如果你看到下面这一行,这意味着 RPC 服务器在你的设备上成功启动。

INFO:root:RPCServer: bind to 0.0.0.0:9090

在本地机器上声明和交叉编译内核

注意

现在我们回到本地机器,它安装了完整的 TVM(含 LLVM)。

这里将在本地机器上声明简单的内核:

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)

然后交叉编译内核。对于 Raspberry Pi 3B 来说,目标应该是 'llvm -mtriple=armv7l-linux-gnueabihf',但在这里使用 'llvm' 来使这个教程可以在网页构建服务器上运行。请看下面区块中的详细说明。

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,并将 target 中的 build 改为适合你设备的目标 triple。对于不同的设备,目标 triple 可能是不同的。例如,对于 Raspberry Pi 3B,它是 'llvm -mtriple=armv7l-linux-gnueabihf',对于 RK3399,它是 'llvm -mtriple=aarch64-linux-gnu'

通常,你可以通过在你的设备上运行 gcc -v 来查询目标,并寻找以 Target: 开始的行。(尽管它可能仍然是松散的配置)。

除了-mtriple,你还可以设置其他的编译选项,比如:

-mcpu=<cpuname>

指定一个当前架构中的特定芯片来生成代码。默认情况下,这是从目标芯片中推断出来的,并自动检测到当前架构。

-mattr=a1,+a2,-a3,…

覆盖或控制目标的特定属性,比如是否启用 SIMD 操作。默认的属性集是由当前的 CPU 设置的。要获得可用的属性列表,你可以这样做:

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

这些选项与 llc 一致。建议将目标 triple 和特性集设置为包含可用的特定特性,这样我们就可以充分利用板子的特性。你可以从 LLVM 的交叉编译指南 中找到更多关于交叉编译属性的细节。

通过 RPC 远程运行 CPU 内核#

展示如何在远程设备上运行已经生成的 CPU 内核。

从远程设备获得 RPC 会话。

if local_demo:
    remote = rpc.LocalSession()
else:
    # 将其改为你的目标设备的 IP 地址
    host = "10.77.1.162"
    port = 9090
    remote = rpc.connect(host, port)

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

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

# create arrays on the remote device
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)
# the function will run on the remote device
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.396e-07 secs/op

通过 RPC 远程运行 OpenCL 内核#

对于远程 OpenCL 设备,工作流程几乎和上面一样。你可以定义内核,上传文件,并通过 RPC 运行。

注意

Raspberry Pi 不支持 OpenCL,以下代码是在 Firefly-RK3399 上测试的。你可以按照这个 教程 来设置操作系统和 RK3399 的 OpenCL 驱动。

此外,还需要在 rk3399 板上启用 OpenCL 来构建运行时。在 TVM 根目录下,执行

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

下面的函数显示了如何远程运行 OpenCL 内核:

def run_opencl():
    # NOTE: This is the setting for my rk3399 board. You need to modify
    # them according to your environment.
    opencl_device_host = "10.77.1.145"
    opencl_device_port = 9090
    target = tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")

    # create schedule for the above "add one" compute declaration
    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)

    # export and upload
    path = temp.relpath("lib_cl.tar")
    func.export_library(path)
    remote.upload(path)
    func = remote.load_module("lib_cl.tar")

    # run
    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 服务器。

  • 设置目标设备配置,在本地机器上交叉编译内核。

  • 通过 RPC API 远程上传和运行内核。