TVM:交叉编译和RPC

10 篇文章 2 订阅
6 篇文章 1 订阅

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 远程上传并运行核
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值