Relay Arm® 计算库集成

Relay Arm® 计算库集成#

Author: Luke Hutton

简介#

Arm 计算库（Arm Compute Library，简称 ACL）是开源项目，为 Arm CPU 和 GPU 提供加速内核。目前，集成将算子卸载到 ACL 以使用库中手工编写的汇编例程。通过将 Relay 图中的选定算子卸载到 ACL，可以在此类设备上实现性能提升。

安装 Arm 计算库#

在安装 Arm 计算库之前，了解要构建的架构非常重要。确定架构的一种方法是使用 lscpu 并查找 CPU 的“型号名称”。然后，您可以通过在线查找来确定架构。

TVM 仅支持单一版本的 ACL，目前为 v21.08，有两种推荐的构建和安装所需库的方法：

  • 使用位于 docker/install/ubuntu_download_arm_compute_lib_binaries.sh 的脚本。您可以使用此脚本下载指定架构和扩展的 ACL 二进制文件，这些文件将安装到 install_path 指定的位置。

  • 或者，您可以从以下网址下载预构建的二进制文件：ARM-software/ComputeLibrary。使用此包时，您需要选择所需架构和扩展的二进制文件，然后确保 CMake 可以找到它们：

    cd <acl-prebuilt-package>/lib
mv ./<architecture-and-extensions-required>/* .

在这两种情况下，您都需要将 USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR 设置为 ACL 包的路径。CMake 将在 /path-to-acl/ 以及 /path-to-acl/lib 和 /path-to-acl/build 中查找所需的二进制文件。有关如何使用这些配置选项的更多信息，请参阅以下部分。

构建支持 ACL 的 TVM#

当前实现在 CMake 中有两个独立的构建选项。这种分离的原因是 ACL 不能在 x86 机器上使用。然而，仍然希望在 x86 机器上编译 ACL 运行时模块。

  • USE_ARM_COMPUTE_LIB=ON/OFF - 启用此标志将添加对编译 ACL 运行时模块的支持。

  • USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR=ON/OFF/path-to-acl - 启用此标志将允许图执行器计算 ACL 卸载的函数。

这些标志可以根据您的设置在不同的场景中使用。例如，如果您想在 x86 机器上编译 ACL 模块，然后通过 RPC 在远程 Arm 设备上运行该模块，您需要在 x86 机器上使用 USE_ARM_COMPUTE_LIB=ON，并在远程 AArch64 设备上使用 USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR=ON。

默认情况下，这两个选项都设置为 OFF。使用 USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR=ON 将意味着 CMake 会在默认位置搜索 ACL 二进制文件（参见 https://cmake.org/cmake/help/v3.4/command/find_library.html）。除此之外，还会搜索 /path-to-tvm-project/acl/。您可能需要设置自己的路径来定位 ACL。这可以通过在 ON 的位置指定路径来完成。

这些标志应在您的 config.cmake 文件中设置。例如：

set(USE_ARM_COMPUTE_LIB ON)
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR /path/to/acl)

用法#

备注

本节内容可能不会随着 API 的变化而保持最新。

创建 Relay 图。这可能是单一算子或整个图。目的是可以输入任何 Relay 图。ACL 集成将仅选择支持的算子进行卸载，而其余部分将通过 TVM 计算。（在此示例中，将使用单个 max_pool2d 算子）。

import tvm
from tvm import relay

data_type = "float32"
data_shape = (1, 14, 14, 512)
strides = (2, 2)
padding = (0, 0, 0, 0)
pool_size = (2, 2)
layout = "NHWC"
output_shape = (1, 7, 7, 512)

data = relay.var('data', shape=data_shape, dtype=data_type)
out = relay.nn.max_pool2d(data, pool_size=pool_size, strides=strides, layout=layout, padding=padding)
module = tvm.IRModule.from_expr(out)

为 ACL 注解和分区图。

from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib
module = partition_for_arm_compute_lib(module)

构建 Relay 图。

target = "llvm -mtriple=aarch64-linux-gnu -mattr=+neon"
with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
    lib = relay.build(module, target=target)

导出模块。

lib_path = '~/lib_acl.so'
cross_compile = 'aarch64-linux-gnu-c++'
lib.export_library(lib_path, cc=cross_compile)

运行推理。这必须在 Arm 设备上进行。如果在 x86 设备上编译并在 AArch64 上运行，请考虑使用 RPC 机制。使用 RPC 机制的教程

dev = tvm.cpu(0)
loaded_lib = tvm.runtime.load_module('lib_acl.so')
gen_module = tvm.contrib.graph_executor.GraphModule(loaded_lib['default'](dev))
d_data = np.random.uniform(0, 1, data_shape).astype(data_type)
map_inputs = {'data': d_data}
gen_module.set_input(**map_inputs)
gen_module.run()

更多例子#

上面的示例仅展示了如何使用 ACL 卸载单个 Maxpool2D 的基本示例。如果您想查看每个已实现算子和网络的更多示例，请参阅测试：tests/python/contrib/test_arm_compute_lib。在这里，您可以修改 test_config.json 以配置如何在 infrastructure.py 中创建远程设备，从而影响运行时测试的运行方式。

test_config.json 的示例配置：

  • connection_type - RPC 连接的类型。选项：local、tracker、remote。

  • host - 要连接的主机设备。

  • port - 连接时使用的端口。

  • target - 用于编译的目标。

  • device_key - 通过跟踪器连接时的设备密钥。

  • cross_compile - 从非 Arm 平台连接时交叉编译器的路径，例如 aarch64-linux-gnu-g++。

{
  "connection_type": "local",
  "host": "127.0.0.1",
  "port": 9090,
  "target": "llvm -mtriple=aarch64-linux-gnu -mattr=+neon",
  "device_key": "",
  "cross_compile": ""
}

算子支持#

Relay 节点

标记

nn.conv2d
fp32:

Simple: nn.conv2d Composite: nn.pad?, nn.conv2d, nn.bias_add?, nn.relu?

Normal and depth-wise (when kernel is 3x3 or 5x5 and strides are 1x1 or 2x2) convolution supported. Grouped convolution is not supported.

qnn.conv2d
uint8:

Composite: nn.pad?, nn.conv2d, nn.bias_add?, nn.relu?, qnn.requantize

Normal and depth-wise (when kernel is 3x3 or 5x5 and strides are 1x1 or 2x2) convolution supported. Grouped convolution is not supported.

nn.dense
fp32:

Simple: nn.dense Composite: nn.dense, nn.bias_add?

qnn.dense
uint8:

Composite: qnn.dense, nn.bias_add?, qnn.requantize

nn.max_pool2d

fp32, uint8

nn.global_max_pool2d

fp32, uint8

nn.avg_pool2d
fp32:

Simple: nn.avg_pool2d

uint8:

Composite: cast(int32), nn.avg_pool2d, cast(uint8)

nn.global_avg_pool2d
fp32:

Simple: nn.global_avg_pool2d

uint8:

Composite: cast(int32), nn.avg_pool2d, cast(uint8)

power(of 2) + nn.avg_pool2d + sqrt

A special case for L2 pooling.

fp32:

Composite: power(of 2), nn.avg_pool2d, sqrt

reshape

fp32, uint8

maximum

fp32

add

fp32

qnn.add

uint8

备注

复合算子是一系列映射到单个 Arm 计算库算子的算子。您可以从 Arm 计算库的角度将其视为单个融合算子。'?' 表示构成复合算子的算子系列中的可选算子。

添加新算子#

添加新算子需要对一系列地方进行更改。本节将提示需要更改的内容和位置，但不会深入探讨单个算子的复杂性。这留给开发人员自行处理。

需要对一系列文件进行更改：

  • python/relay/op/contrib/arm_compute_lib.py 在此文件中，使用 op.register 装饰器定义希望卸载的算子。这将意味着注释传递将此算子识别为可 ACL 卸载的。

  • src/relay/backend/contrib/arm_compute_lib/codegen.cc 实现 Create[OpName]JSONNode 方法。这是声明算子应如何由 JSON 表示的地方。这将用于创建 ACL 模块。

  • src/runtime/contrib/arm_compute_lib/acl_runtime.cc 实现 Create[OpName]Layer 方法。这是定义如何使用 JSON 表示创建 ACL 函数的地方。只需定义如何从 JSON 表示转换为 ACL API。

  • tests/python/contrib/test_arm_compute_lib 为给定算子添加单元测试。