远程分析和测试基于 TVM RPC 的移动电话深度学习交叉编译

•


TVM 堆栈是一个端到端的编译堆栈，用于将深度学习工作负载部署到所有硬件后端。 感谢 TVM 堆栈的 NNVM 编译器支持，我们现在可以直接编译来自深度学习框架的描述，并将它们编译成裸机代码。 TVM 令人印象深刻的特性之一是它能够在不同平台（如 GPU 和移动电话（将支持更多硬件后端））上部署计算工作负载。

然而，当我们想要测试和分析交叉编译时，很难在树莓派或移动电话等异构设备上测试不同的计算工作负载。 为了优化计算任务，必须在开发 PC 上编辑代码，编译，部署到设备，测试，然后再次修改代码以查看是否加速。 工作流程如下：

image

有没有办法加快这个过程？

今天，我们介绍一种在 Android 手机上部署和测试 TVM 工作负载的方法。 我们为 Java 开发了一个 TVM 运行时，并在此基础上构建了一个 Android APP。 Android APP 将共享库作为输入，并在移动电话上运行编译后的函数。 因此，我们的工作流程简化为：

image

借助 TVM RPC，可以在远程设备上构建 TVM 函数和 NDArray。 交叉编译到不同平台的能力使得在一个平台上开发并在另一个平台上测试变得容易。

该过程如下图所示

image

在 Android 手机上运行 TVM APP

您可以在 apps/android_rpc 中找到 Android RPC APP。 请按照说明为您的 Android 设备构建。 APK 构建完成后，使用 apps/android_rpc/dev_tools 对其进行签名并安装在手机上。 该 APP 看起来像

image image

通常我们无法在移动电话上启动独立服务器，而是启动代理服务器并使用我们的应用程序连接。

python -m tvm.exec.rpc_proxy

在手机上创建 NDArray

现在我们可以从笔记本电脑连接到代理服务器

from tvm.contrib import rpc
remote = rpc.connect("0.0.0.0", 9090, key="android")

这将为我们提供一个处理程序 remote，我们可以使用它与移动电话进行通信。 例如，以下行在手机的 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 手机的架构。 我们在三星 Galaxy S6 Edge 上进行了测试，它有一个 Mali-T760 GPU。 这是这款手机的 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 Function 和 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

这将在您的本地 Maven 仓库中编译、打包和安装 tvm4j。 请参阅 tvm4j 以获取更多信息。

在 iPhone/iPad 上进行远程分析和测试

除了 Android RPC 应用程序外，我们还提供了一个 iOS RPC 应用程序，通过它可以轻松地在 iPhone 或 iPad 上分析和测试 TVM 计算工作负载。 它的工作方式与 Android 上的几乎相同，但需要 XCode 和 iOS 设备。