Cross Compilation and RPC

本文将介绍交叉编译以及TVM通过RPC在远程设备上执行。

通过交叉编译和RPC，可以在本地机器上编译程序，然后在远程设备上运行。这在远程设备的资源有限时十分重要，例如Rasberry Pi和移动平台，本文将介绍Rasberry Pi的CPU例子和Firefly-RK3399的OpenCL例子

在设备上构建TVM Runtime

第一步是在远程设备上构建TVM runtime

因为我们在本地机器上编译，远程设备只用来执行构建好的代码，所以只需要在远程设备上构建TVM runtime。

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

成功构建runtime后，需要设置环境变量，修改~/.bashrc，假定TVM的目录是~/tvm

export PYTHONPATH=$PYTHONPATH:~/tvm/python

更新环境变量

source ~/.bashrc

在设备上启动RPC服务器

在远程设备运行代码，例如Rasberry 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

在本地机器上定义并交叉编译kernel

这里使用有LLVM的TVM，首先在本地机器定义一个简单的kernel

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)

然后交叉编译，目标应该是llvm -mtriple=armv7l-linux-gnueabihf，这是Rasberry Pi 3B的选项，不过我们这里使用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")
# 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为对应的设备，例如Rasberry Pi 3B使用llvm -mtriple=armv7l-linux-gnueabihf，RK3399使用llvm -mtriple=aarch64-linux-gnu。

通常，可以通过在目标设备上运行gcc -v查看Target

通过RPC远程启动CPU kernel

首先从远程设备获取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)

上传lib到远程设备，然后调用本地编译器重新链接它们，现在func是一个远程的module对象

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)

当需要评估远程设备上kernel的性能时，需要避免提前网络的影响，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.187e-07 secs/op

通过RPC远程启动OpenCL Kernel

对于远程OpenCL设备，工作流大部分都是相同的，可以定义kernel，上传文件，然后通过RPC运行

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

下面的函数展示如何远程运行OpenCL kernel

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!")