交叉编译和 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")
mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "add_one"))

然后我们交叉编译内核。对于 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.compile(mod, target=target)
# save the lib at a local temp folder
temp = utils.tempdir()
path = temp.relpath("lib.tar")
func.export_library(path)

注意

要使用真正的远程设备运行本教程，请将 local_demo 更改为 False，并将 build 中的 target 替换为适合您设备的正确目标三元组。不同设备的目标三元组可能有所不同。例如，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 一致。建议设置目标三元组和特性集以包含可用的特定特性，以便我们可以充分利用板卡的特性。您可以从 LLVM 交叉编译指南 中找到有关交叉编译属性的更多详细信息。

通过 RPC 远程运行 CPU 内核

我们展示如何在远程设备上运行生成的 CPU 内核。首先，我们从远程设备获取 RPC 会话。

if local_demo:
    remote = rpc.LocalSession()
else:
    # The following is my environment, change this to the IP address of your target device
    host = "10.77.1.162"
    port = 9090
    remote = rpc.connect(host, port)

将库上传到远程设备，然后调用设备本地编译器重新链接它们。现在 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.228e-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
    mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]))
    sch = tvm.tir.Schedule(mod)
    (x,) = sch.get_loops(block=sch.get_block("B"))
    xo, xi = sch.split(x, [None, 32])
    sch.bind(xo, "blockIdx.x")
    sch.bind(xi, "threadIdx.x")
    func = tvm.compile(sch.mod, 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 远程上传和运行内核。