以下内容翻译自:Remote Profile and Test Deep Learning Cross Compilation on Mobile Phones with TVM RPC
TVM 堆栈是端到端的编译堆栈,可将深度学习工作负载部署到所有硬件后端。由于 NNVM 编译器支持 TVM 堆栈,我们现在可以直接编译来自深度学习框架的描述并生成裸机代码。TVM 一个令人印象深刻的特性是它能够在不同的平台上部署计算工作负载,比如 GPU 和手机(将支持更多的硬件后端)。
然而,当我们想要测试和剖析交叉编译时,很难在异构设备(如树莓派或手机)上测试不同的计算工作负载。为了优化计算任务,必须在开发 PC 上 编辑代码,编译,部署到设备,测试,然后再次修改代码以查看是否加速。工作流程看起来像:
有什么办法可以加快这一过程吗?
今天我们介绍一种在 Android 手机上部署和测试 TVM 工作负载的方法。我们为 Java 开发了一个 TVM 运行库,并在其上构建了一个 Android 应用程序。Android APP 将共享库作为输入,并在手机上运行编译后的函数。因此我们的工作流程简化为:
借助 TVM RPC,我们可以在远程设备上构建 TVM 函数和 NDArray。交叉编译到不同平台的能力使得在一个平台上开发并在另一个平台上进行测试变得容易。
该过程如下所示:
Android__TVM_APP_18">在 Android 手机上运行 TVM APP
你可以在 apps/android_rpc 中找到 Android RPC APP。请按照说明为您的 Android 设备构建。一旦生成了 APK,请使用 apps/android_rpc/dev_tools 对其进行签名并将其安装在手机上。该 APP 看起来像:
通常,我们无法在手机上启动独立服务器,为此启动代理服务器并使用我们的应用进行连接。
python -m tvm.exec.rpc_proxy
在手机上创建 NDArray
现在我们可以从笔记本电脑连接到代理服务器:
from tvm.contrib import rpc
remote = rpc.connect("0.0.0.0", 9090, key="android")
这将给我们一个可以用与手机通信的远程句柄。例如,以下几行在手机的 GPU上 创建了一个1024x1024矩阵:
A = tvm.nd.array(
np.random.uniform(size=(1024, 1024)).astype(dtype),
ctx = remote.cl(0))
当笔记本电脑调用A.asnumpy()
时,矩阵A
将被复制到手机的 RAM,然后通过代理服务器传输到笔记本电脑。TVM RPC 接口对用户是透明的。
手机上的 GEMM(矩阵乘法)
现在我们将介绍如何在 Android 手机上测试矩阵乘法。首先让我们定义非常简单的 GEMM 方案:
import tvm
def gemm(N, bn):
A = tvm.placeholder((N, N), name='A')
B = tvm.placeholder((N, N), name='B')
k = tvm.reduce_axis((0, N), name='k')
C = tvm.compute(
(N, N),
lambda ii, jj: tvm.sum(A[ii, k] * B[k, jj], axis=k),
name='C')
s = tvm.create_schedule(C.op)
block_x = tvm.thread_axis("blockIdx.x")
thread_x = tvm.thread_axis("threadIdx.x")
bo, bi = s[C].split(C.op.axis[0], factor=bn)
to, ti = s[C].split(C.op.axis[1], factor=bn)
s[C].bind(bi, block_x)
s[C].bind(ti, thread_x)
print(tvm.lower(s, [A, B, C], simple_mode=True))
return tvm.build(s, [A, B, C],
"opencl",
target_host="llvm -target=arm64-linux-android",
name="gemm_gpu")
除了最后一行,没有什么特别的。这里我们将目标设置为opencl
,因为这是 Mali GPU 支持的计算语言。请注意,我们将target_host
设置为llvm -target = arm64-linux-android
,这取决于您 Android 手机的架构。我们在配备 Mali-T760 GPU 的三星 Galaxy S6 Edge 上进行测试。这是这款手机的 CPU 信息:
$ adb shell
shell@zenltechn:/ $ cat /proc/cpuinfo
Processor : AArch64 Processor rev 2 (aarch64)
processor : 0
processor : 1
processor : 2
processor : 3
processor : 4
processor : 5
processor : 6
processor : 7
Features : fp asimd aes pmull sha1 sha2 crc32
CPU implementer : 0x41
CPU architecture: AArch64
CPU variant : 0x0
CPU part : 0xd03
CPU revision : 2
Hardware : SAMSUNG Exynos7420
请参阅 target triple 了解 LLVM 的编译选项。我们使用 tvm.contrib.ndk 为 Android 系统创建共享库:
from tvm.contrib import rpc, util, ndk
N = 1024
f = gemm(N, bn = 256)
temp = util.tempdir()
path_dso = temp.relpath("gemm_gpu.so")
f.export_library(path_dso, ndk.create_shared)
ndk.create_shared 读取环境变量TVM_NDK_CC
以查找 Android 设备的编译器和链接器。我们可以很容易地使用 NDK 为我们的设备生成独立的工具链。例如,以下命令为 ARM64 Android 设备生成独立编译器和连接器:
cd /opt/android-ndk/build/tools/
./make-standalone-toolchain.sh --platform=android-24 --use-llvm --arch=arm64 --install-dir=/opt/android-toolchain-arm64
如果一切顺利,我们会有一个“gemm_gpu.so”共享库。现在让我们将其上传到手机,使手机加载模块并获得远程句柄:
remote = rpc.connect("0.0.0.0", 9090, key="android")
remote.upload(path_dso)
f = remote.load_module("gemm_gpu.so")
创建远程阵列并打印运行时间:
ctx = remote.cl(0)
import numpy as np
a_np = np.random.uniform(size=(N, N)).astype("float32")
b_np = np.random.uniform(size=(N, N)).astype("float32")
a = tvm.nd.array(a_np, ctx)
b = tvm.nd.array(b_np, ctx)
c = tvm.nd.array(np.zeros((N, N), dtype="float32"), ctx)
time_f = f.time_evaluator(f.entry_name, ctx, number=5)
cost = time_f(a, b, c).mean
print('%g secs/op, %g GFLOPS' % (cost, ngflops(N) / cost))
现在我们可以在 PC 上验证结果:
np.testing.assert_almost_equal(
c.asnumpy(),
a_np.dot(b_np),
decimal=3)
在上面的情况下,我们开发并交叉编译为手机的二进制文件。通过代理服务器,上传二进制文件到手机并运行在其 JVM 中。这种方法使得在 Android 上开发和测试不同的计算工作量变得很容易。
TVM 的 Java 运行时
Android APP 建立在 Java 运行时之上,它为 TVM 功能和 NDArray 提供了最低限度的支持。这是一个在 TVM4J 中注册函数的例子:
Function func = Function.convertFunc(new Function.Callback() {
@Override public Object invoke(TVMValue... args) {
StringBuilder res = new StringBuilder();
for (TVMValue arg : args) {
res.append(arg.asString());
}
return res.toString();
}
});
TVMValue res = func.pushArg("Hello").pushArg(" ").pushArg("World!").invoke();
assertEquals("Hello World!", res.asString());
res.release();
func.release();
正如我们在 GEMM 部分看到的,可以通过 Python 构建共享库并通过 Java 来执行:
import ml.dmlc.tvm.Module;
import ml.dmlc.tvm.NDArray;
import ml.dmlc.tvm.TVMContext;
import java.io.File;
import java.util.Arrays;
public class LoadAddFunc {
public static void main(String[] args) {
String loadingDir = args[0];
Module fadd = Module.load(loadingDir + File.separator + "add_cpu.so");
TVMContext ctx = TVMContext.cpu();
long[] shape = new long[]{2};
NDArray arr = NDArray.empty(shape, ctx);
arr.copyFrom(new float[]{3f, 4f});
NDArray res = NDArray.empty(shape, ctx);
fadd.entryFunc().pushArg(arr).pushArg(arr).pushArg(res).invoke();
System.out.println(Arrays.toString(res.asFloatArray()));
arr.release();
res.release();
fadd.release();
}
}
按照“安装指南”构建 TVM 库后,运行
make jvmpkg
make jvminstall
这将编译、打包并安装 tvm4j 到您的本地 maven 仓库。 请参阅 TVM4J 了解更多信息。
iPhone/iPad 上的远程配置和测试
除了 Android RPC 应用程序外,我们还提供 iOS RPC app,通过它我们可以轻松地在 iPhone 或 iPad 上分析和测试 TVM 计算工作负载。它的工作原理几乎与 Android 相同,而 XCode 和 iOS 设备是必需的。