交叉编译与RPC#

参考:tutorial-cross-compilation-and-rpc

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

通过交叉编译和 RPC,您可以 在本地机器上编译程序,然后在远程设备上运行它。当远程设备资源有限时,例如树莓派和移动平台,这一方法非常有用。在本教程中,将使用树莓派作为 CPU 示例,以及 Firefly-RK3399 作为 OpenCL 示例。

在设备上构建 TVM 运行时#

首先,需要在远程设备上构建 TVM 运行时。

备注

本节和下一节中的所有指令都应该在目标设备上执行,例如树莓派。假设目标是运行 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 服务器#

要在远程设备(本例中为 Raspberry Pi)上启动 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

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

备注

现在回到本地机器,该机器已完全安装了 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")
# 将库保存到本地临时文件夹
temp = utils.tempdir()
path = temp.relpath("lib.tar")
func.export_library(path)

备注

请注意,为了在真实远程设备上运行本教程,请将 local_demo 设置为False,并将 build 中的 target 替换为您设备的适当目标三元组。不同设备的目标三元组可能有所不同。例如,对于树莓派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=<您的设备目标三元组> -mattr=help

这些选项与llc保持一致。建议设置目标三元组和功能集以包含特定的可用功能,这样我们就可以充分利用板卡的功能。

您可以从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)

将库上传到远程设备,然后调用设备本地的编译器重新链接它们。现在 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(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
2024-11-26 15:27:49.826 INFO load_module /tmp/tmplbfhxmbe/lib.tar

当你想评估远程设备上内核的性能时,重要的是要避免网络开销。time_evaluator 将返回远程函数,该函数在远程设备上多次运行该函数,测量每次运行的成本,并返回测量到的成本。网络开销被排除在外。

time_f = func.time_evaluator(func.entry_name, dev, number=10)
cost = time_f(a, b).mean
print(f"{cost:g} secs/op")
3.097e-07 secs/op

通过 RPC 远程运行 OpenCL 内核#

对于远程 OpenCL 设备,工作流程几乎与上述相同。

您可以定义内核,上传文件,并通过 RPC 运行。

备注

树莓派不支持 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 内核:

Hide code cell content
def run_opencl(opencl_device_host = "10.77.1.145"):
    # NOTE: 这是我的rk3399开发板的设置。您需要根据您的环境进行相应的修改。
    opencl_device_port = 9090
    target = tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")

    # 为上述 "add one" 计算声明创建调度
    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 上传并远程运行内核。