TVM：交叉編譯和RPC

之前我們介紹了 TVM 的安裝、本機(jī)demo和樹(shù)莓派遠(yuǎn)程demo。本文將介紹了在 TVM 中使用 RPC 進(jìn)行交叉編譯和遠(yuǎn)程設(shè)備執(zhí)行。

通過(guò)交叉編譯和 RPC，我們可以在本地機(jī)器上編譯程序，然后在遠(yuǎn)程設(shè)備上運(yùn)行它。 當(dāng)遠(yuǎn)程設(shè)備資源有限時(shí)很有用，例如 Raspberry Pi 和移動(dòng)平臺(tái)。 在本文中，我們將使用 Raspberry Pi 作為 CPU 示例，使用 Firefly-RK3399 作為 OpenCL 示例。

在遠(yuǎn)程設(shè)備上構(gòu)建 TVM Runtime

首先我們要在遠(yuǎn)程設(shè)備上編譯安裝 TVM Runtime，注意這里我們對(duì)模型的編譯是在本機(jī)進(jìn)行的，而遠(yuǎn)程設(shè)備只需要運(yùn)行模型即可，因此只需要構(gòu)建 TVM Runtime。

注意：本節(jié)和下一節(jié)中的所有指令都應(yīng)在目標(biāo)設(shè)備上執(zhí)行，例如樹(shù)莓派。 我們假設(shè)它運(yùn)行著 Linux。

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

將 Python 路徑添加到環(huán)境變量：

export PYTHONPATH=$PYTHONPATH:/path/to/tvm/python

在遠(yuǎn)程設(shè)備上設(shè)置 RPC 服務(wù)器

在遠(yuǎn)程設(shè)備（如本例中的樹(shù)莓派）上運(yùn)行以下命令來(lái)開(kāi)啟 RPC 服務(wù)器：

python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090

如果看到下面這行說(shuō)明遠(yuǎn)程設(shè)備上的 RPC 服務(wù)已經(jīng)成功開(kāi)啟了：

INFO:root:RPCServer: bind to 0.0.0.0:9090

在本機(jī)上聲明并交叉編譯核

注意：現(xiàn)在我們回到本地機(jī)器了，之后的操作都是在本機(jī)（含有完整的，帶有 LLVM 的 TVM）上進(jìn)行。

我們現(xiàn)在本機(jī)上聲明一個(gè)簡(jiǎn)單的核：

import numpy as npimport tvm from tvm import te from tvm import rpc from tvm.contrib import utilsn = 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)

然后我們來(lái)對(duì)核進(jìn)行交叉編譯。對(duì)于樹(shù)莓派3B來(lái)說(shuō)，target 應(yīng)該是 llvm -mtriple=armv7l-linux-gnueabihf’ 。如果真的有一個(gè)遠(yuǎn)程設(shè)備樹(shù)莓派的話可以將下面的 local_demo 改為 False ，否則還是保留為 True 使得本 demo 可以正常運(yùn)行。

local_demo = Trueif local_demo:target = "llvm" else:target = "llvm -mtriple=armv7l-linux-gnueabihf"func = tvm.build(s, [A, B], target=target, name="add_one") # 在本地的臨時(shí)目錄下保存一個(gè) lib temp = utils.tempdir() path = temp.relpath("lib.tar") func.export_library(path)

注意：要使用真正的遠(yuǎn)程設(shè)備運(yùn)行本教程，請(qǐng)將 local_demo 更改為 False 并將 build 中的 target 替換為適合我們?cè)O(shè)備的目標(biāo)三元組（target triple）。不同設(shè)備的（target triple）可能不同。例如，對(duì)于 Raspberry Pi 3B，它是'llvm -mtriple=armv7l-linux-gnueabihf'，對(duì)于 RK3399，它是'llvm -mtriple=aarch64-linux-gnu'。

通常，我們可以通過(guò)在您的設(shè)備上運(yùn)行 gcc -v 來(lái)查詢目標(biāo)，并查找以 Target: 開(kāi)頭的行（盡管它可能仍然是一個(gè)寬松的配置。）

除了 -mtriple，您還可以設(shè)置其他編譯選項(xiàng)，例如：

• -mcpu=<cpuname>
  指定當(dāng)前架構(gòu)中的特定芯片以為其生成代碼。默認(rèn)情況下，這是從 target triple 推斷出來(lái)的，并自動(dòng)檢測(cè)到當(dāng)前架構(gòu)。

• -mattr=a1,+a2,-a3,…
  覆蓋或控制目標(biāo)的特定屬性，例如是否啟用 SIMD 操作。默認(rèn)屬性集由當(dāng)前 CPU 設(shè)置。要獲取可用屬性列表，我們可以執(zhí)行以下操作：

llc -mtriple=\<your device target triple\> -mattr=help

這些選項(xiàng)與 llc 一致。建議將 target triple和功能集設(shè)置為包含可用的特定功能，以便我們可以充分利用板的功能。可以到 LLVM 交叉編譯指南中找到有關(guān)交叉編譯屬性的更多詳細(xì)信息。

通過(guò)RPC在遠(yuǎn)程運(yùn)行CPU核

接下來(lái)是如何將生成的CPU核運(yùn)行在遠(yuǎn)程設(shè)備上，首先我們建立與遠(yuǎn)程設(shè)備的 RPC 會(huì)話。

if local_demo:remote = rpc.LocalSession() else:# 下面的IP是筆者的，請(qǐng)大家換成自己的遠(yuǎn)程設(shè)備的IPhost = "10.206.105.111"port = 9090remote = rpc.connect(host, port)

將 lib 上傳到遠(yuǎn)程設(shè)備，然后調(diào)用設(shè)備本地編譯器來(lái)重新鏈接它們。 現(xiàn)在 func 是一個(gè)遠(yuǎn)程模塊對(duì)象。

remote.upload(path) func = remote.load_module("lib.tar")# 在遠(yuǎn)程設(shè)備上創(chuàng)建數(shù)組 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 將會(huì)運(yùn)行在遠(yuǎn)程設(shè)備上 func(a, b) np.testing.assert_equal(b.numpy(), a.numpy() + 1)

如果我們想要在遠(yuǎn)程設(shè)備上評(píng)估核的性能時(shí)，我們需要避免網(wǎng)絡(luò)開(kāi)銷。 time_evaluator 將返回一個(gè)遠(yuǎn)程函數(shù)，該函數(shù)多次運(yùn)行該函數(shù)，測(cè)量遠(yuǎn)程設(shè)備上每次運(yùn)行的成本并返回測(cè)量的成本。并將網(wǎng)絡(luò)開(kāi)銷排除在外。

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

通過(guò)RPC在遠(yuǎn)程設(shè)備上運(yùn)行OpenCL核

對(duì)于遠(yuǎn)程 OpenCL 設(shè)備，整個(gè)流程和上面幾乎是一樣的。我們定義自己的和核，上傳文件，并通過(guò) RPC 運(yùn)行。

注意：樹(shù)莓派并不支持 OpenCL，以下代碼是在 Firefly-RK3399 上進(jìn)行測(cè)試的。大家可以通過(guò)這個(gè)教程來(lái)為 RK3399 配置操作系統(tǒng)和 OpenCL。

同樣我們需要再 RK3399 上構(gòu)建 TVM Runtiime（注意要在 config.cmake 中啟用 OpenCL），在 tvm 根目錄下，執(zhí)行：

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

接下來(lái)。我們通過(guò)以下代碼來(lái)遠(yuǎn)程運(yùn)行 OpenCL 核：

def run_opencl():# 注意，這里是我自己的 RK3399 板子的設(shè)置，你可以根據(jù)自己的環(huán)境進(jìn)行調(diào)整opencl_device_host = "10.77.1.145"opencl_device_port = 9090target = tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")# 為上述 'add one' 計(jì)算聲明創(chuàng)建 schedules = 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)# 導(dǎo)出并上傳path = temp.relpath("lib_cl.tar")func.export_library(path)remote.upload(path)func = remote.load_module("lib_cl.tar")# 運(yùn)行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!")

總結(jié)

本文提供了 TVM 中交叉編譯和 RPC 功能的演示。

• 在遠(yuǎn)程設(shè)備上設(shè)置 RPC 服務(wù)器
• 設(shè)置設(shè)備的 target配置 并在本機(jī)上交叉編譯核
• 通過(guò) RPC 遠(yuǎn)程上傳并運(yùn)行核

總結(jié)

以上是生活随笔為你收集整理的TVM：交叉编译和RPC的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。