交叉編譯和 RPC

超神經HyperAI發表於2023-03-09

本篇文章譯自英文檔案 Cross Compilation and RPC 作者是 Ziheng Jiang,Lianmin Zheng。更多 TVM 中文檔案可訪問 →TVM 中文站

本教程介紹瞭如何在 TVM 中使用 RPC 進行交叉編譯和遠端裝置執行。

利用交叉編譯和 RPC,可以實現程式在本地機器編譯,在遠端裝置執行。這個特性在遠端裝置資源有限時(如在樹莓派和移動平臺上)尤其有用。本教程將把樹莓派作為 CPU 示例,把 Firefly-RK3399 作為 OpenCL 示例進行演示。

在裝置上構建 TVM Runtime

首先在遠端裝置上構建 TVM runtime。

注意 本節和下一節中的所有命令都應在目標裝置(例如樹莓派)上執行。假設目標裝置執行 Linux 系統。

由於在本地機器上只做編譯,而遠端裝置用於執行生成的程式碼。所以只需在遠端裝置上構建 TVM runtime。

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

成功構建 runtime 後,要在 ~/.bashrc 檔案中設定環境變數。可以用 vi ~/.bashrc命令編輯 ~/.bashrc,在這個檔案裡新增下面這行程式碼(假設 TVM 目錄在 ~/tvm 中):

export PYTHONPATH=$PYTHONPATH:~/tvm/python

執行 source ~/.bashrc 來更新環境變數。

在裝置上設定 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”,但這裡用的是“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")
# 將 lib 儲存在本地臨時資料夾
temp = utils.tempdir()
path = temp.relpath("lib.tar")
func.export_library(path)

備註

要使本教程執行在真正的遠端裝置上,需要將 local_demo 改為 False,並將 build 中的 target 替換為適合裝置的 target 三元組。不同裝置的 target 三元組可能不同。例如,對於樹莓派 3B,它是 llvm -mtriple=armv7l-linux-gnueabihf;對於 RK3399,它是 llvm -mtriple=aarch64-linux-gnu。

通常,可以在裝置上執行 gcc -v 來查詢 target,尋找以 Target 開頭的行:(儘管它可能仍然是一個鬆散的配置。)
除了 -mtriple,還可設定其他編譯選項,例如:

  • -mcpu=< cpuname>

指定生成的程式碼執行的晶片架構。預設情況這是從 target 三元組推斷出來的,並自動檢測到當前架構。

  • -mattr=a1,+a2,-a3,…

覆蓋或控制 target 的指定屬性,例如是否啟用 SIMD 操作。預設屬性集由當前 CPU 設定。要獲取可用屬性列表,執行:

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

這些選項與 llc 一致。建議設定 target 三元組和功能集,使其包含可用的特定功能,這樣我們可以充分利用單板的功能。檢視 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)

將 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(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)

要想評估核心在遠端裝置上的效能,避免網路開銷很重要。time_evaluator 返回一個遠端函式,這個遠端函式多次執行 func 函式,並測試每一次在遠端裝置上執行的成本,然後返回測試的成本(不包括網路開銷)。

time_f = func.time_evaluator(func.entry_name, dev, number=10)
cost = time_f(a, b).mean
print("%g secs/op" % cost)

輸出結果:

1.369e-07 secs/op

透過 RPC 遠端執行 OpenCL 核心

遠端 OpenCL 裝置的工作流程與上述內容基本相同。可以定義核心、上傳檔案,然後透過 RPC 執行。

備註
樹莓派不支援 OpenCL,下面的程式碼是在 Firefly-RK3399 上測試的。可以按照 教程 為 RK3399 設定 OS 及 OpenCL 驅動程式。
在 rk3399 板上構建 runtime 也需啟用 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 伺服器。
  • 設定目標裝置配置,使得可在本地機器上交叉編譯核心。
  • 透過 RPC API 遠端上傳和執行核心。

下載 Python 原始碼
下載 Jupyter Notebook

以上就是該檔案的全部內容,檢視更多 TVM 中文檔案,請訪問→TVM 中文站

相關文章