交叉编译与RPC

交叉编译与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 set_env

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1732605760.942433  657022 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1732605761.013284  657022 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

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 内核：

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 上传并远程运行内核。