在x86的Docker中构建TVM的ARM环境

前言

  本篇文章介绍如何在x86的docker中构建tvm的ARM环境，以及如何使用RPC使编译在x86的环境中运行在arm的环境中，还介绍了如何在arm环境中进行编译和运行，并提供详细的示例以供验证，其中包括rpc测试代码，acl测试代码，pytorch模型在arm上的推理以及在arm上进行autotvm。
  如下图所示，显示的是x86架构的cpu信息：

  强烈建议使用ubuntu:20.04这个版本，ubuntu:18.04这个版本在升级glibc时会掉坑里！！！

1. 加载arm-ubuntu镜像

  从docker镜像库中拉取arm-ubuntu镜像：

docker pull arm64v8/ubuntu:20.04

  由于本地的cpu是x86架构，没办法直接运行arm架构的镜像，需要借助第三方工具：QEMU
  QEMU是一个通用的开源的跨平台仿真模拟器，可以模拟在特定的体系结构下的应用的执行或者构建，比如在x86的体系结构的操作系统上运行ARM的应用。
  目前看到使用qemu进行模拟的有两种方式：一种是结合docker使用[本博客使用的]，另一种是使用qemu官方的源码进行编译，手动安装相应系统的iso，可以参考这篇博客。

  使用docker构建的arm环境，用lscpu指令查看cpu时，cpu的model name仍然是intel的，arch是aarch64；使用源码编译成的qemu-system-aarch64构建模拟器时会指定具体的cpu型号，比如qemu-system-aarch64 -cpu cortex-a72，由于没有尝试这种方式，因此不确定cpu的model name是intel的，还是arm的。

docker run --rm --privileged multiarch/qemu-user-static --reset -p yes

  这条指令会安装qemu-user-static，安装完毕后就可以正常在x86的docker中运行ARM架构下的ubuntu镜像了，不过这条指令还没有测试。我使用的是下面的方法，可以参考这篇文章：
  先下载qemu-aarch64-static安装包：

  对qemu-aarch64-static进行配置：

sudo cp qemu-aarch64-static /usr/bin/
sudo chmod +x /usr/bin/qemu-aarch64-static
# 注册QEMU虚拟机
docker run --rm --privileged multiarch/qemu-user-static:register

  然后就可以正常的加载arm-ubuntu了：

docker run --platform linux/arm64/v8 -it -v /home/liyanpeng/arm64v8_work:/home/liyanpeng/arm64v8_work -w /home/liyanpeng arm64v8/ubuntu:20.04 bash

# uname -a
# lscpu

2. 安装acl库

  ARM 计算库(Arm Compute Library, ACL)，是为 ARM 架构的 CPU 和 GPU 提供加速内核的开源项目。可以从 ARM-software 下载预构建的二进制文件：

# 将压缩包解压到 acl_tmp 目录
tar -zxvf arm_compute-v22.08-bin-linux-arm64-v8.2-a-neon.tar.gz -C acl_tmp

  直接编译ARM架构下的runtime时会报错，需要手动调整相应的目录，可以参考tvm官方给的一个脚本：ubuntu_download_arm_compute_lib_binaries.sh

cp -r acl_tmp/include acl/
cp -r acl_tmp/arm_compute acl/include/
cp -r acl_tmp/support acl/include/
cp -r acl_tmp/utils acl/include/
cp -r acl_tmp/lib/arm64-v8.2-a-neon acl/lib

3. 编译arm运行时

  在编译之前，仍然需要在arm-ubuntu配置基本的环境，包括C/C++、CMake和Python等基本环境，具体可以参考我的这篇文章：《tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境》，这里不在赘述。

  修改build/config.cmake文件：

set(USE_LLVM OFF)	# line 136(default)
set(USE_ARM_COMPUTE_LIB OFF)	# line 236(default)
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR "/home/liyanpeng/arm64v8_work/acl")	# acl的路径

  进行编译：

cd build
cmake ..
make runtime -j6

  编译成功后的信息如下：

  不要忘了添加tvm的python环境：

export PYTHONPATH=$PYTHONPATH:/home/liyanpeng/arm64v8_work/tvm_work/tvm/python

  tvm版本验证：

import tvm

print(tvm.__version__)

4. 编译在x86运行在arm

4.1 在x86的环境中构建arm的编译环境

  修改build/config.cmake文件：

set(USE_LLVM ON)	# line 136
set(USE_ARM_COMPUTE_LIB ON)	# line 236
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR OFF) # line 237

  进行编译：

cd build
cmake ..
make -j6

  很快就编译好了：

  构建完成后即可在x86的环境中编译arm支持的算子，仅仅是编译，在x86上是不能运行直接运行的。

4.2 测试x86-ubuntu与arm-ubuntu能否ping通

  可以借助RPC(Remote Produce Call)来实现 编译在x86，运行在ARM，因此，需要知道arm-ubuntu的ip地址：
  在arm-ubuntu中安装网络工具包：

apt-get update
# ifconfig
apt-get install net-tools
# ping
apt-get install inetutils-ping

  查看arm-ubuntu的ip地址：

  查看x86-ubuntu的ip地址：

  测试x86-ubuntu与arm-ubuntu能否ping通：

# x86-ubuntu
ping 172.17.0.2

# arm-ubuntu
ping 172.17.0.3

4.3 调用RPC

  在arm-ubuntu环境中启动RPC：

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

  启动成功的信息如下：

  在x86-ubuntu环境中创建rpc_test.py文件，内容如下：

# rpc_test.py
import numpy as np

import tvm
from tvm import te
from tvm import rpc
from tvm.contrib import utils, tar


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)

local_demo = False

if local_demo:
    target = "llvm"
else:
    # target = "llvm -mtriple=armv7l-linux-gnueabihf"     # Raspberry Pi 3B
    # target = "llvm -mtriple=aarch64-linux-gnu"
    # target = tvm.target.arm_cpu() # error: error adding symbols: file in wrong format
    target = "llvm -mtriple=aarch64-linux-gnu -mattr=+neon"

func = tvm.build(s, [A, B], target=target, name="add_one")

# save the lib at a local temp folder
temp = utils.tempdir()
path = temp.relpath("lib_rpc_test.tar")
func.export_library(path, tar.tar)

print("lib path: ", path)

if local_demo:
    remote = rpc.LocalSession()
else:
    # The following is my environment, change this to the IP address of your target device
    host = "172.17.0.5"		# arm-ubuntu ip
    port = 9090
    remote = rpc.connect(host, port)

remote.upload(path)
func = remote.load_module("lib_rpc_test.tar")

# create arrays on the remote device
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)

# the function will run on the remote device
func(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)

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

  上面代码演示了一个加法运算，执行结果如下：

  在arm-ubuntu可以看到来自x86-ubuntu的连接信息：

4.4 ACL的使用

  ACL的使用可以参看tvm官方的示例文档，这里给出了例子，使用方式同上小节的一样用的是RPC：

# acl_test.py
import tvm
from tvm import relay
from tvm import rpc
from tvm.contrib import utils, tar
from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib

import numpy as np


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)

# use a single max_pool2d operator
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)

# annotate and partition the graph for ACL
module = partition_for_arm_compute_lib(module)

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

# export the module
lib_path = './lib_acl.tar'
# cross_compile = 'aarch64-linux-gnu-c++'
# lib.export_library(lib_path, cc=cross_compile)
lib.export_library(lib_path)

# rpc
host = "172.17.0.2"		# arm-ubuntu ip
port = 9090
remote = rpc.connect(host, port)

remote.upload(lib_path)
loaded_lib = remote.load_module("lib_acl.tar")

# run Inference
# dev = tvm.cpu(0)
# loaded_lib = tvm.runtime.load_module('lib_acl.so')
dev = remote.cpu(0)
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}
module.set_input(**map_inputs)
module.run()

# get output
output = module.get_output(0)
print("TVM MaxPool2d[acl] output: ", output)

  运行结果如下：

  上面的示例仅显示了ACL如何用于单个Maxpool2D的基本示例。如果想看到网络中每个算子的实现，请参阅：tests/python/contrib/test_arm_compute_lib。

5. arm版的tvm编译和运行时环境

5.1 构建arm版的tvm编译和运行时环境

  ARM版的tvm编译和运行时环境的构建同x86版的几乎一样，可以按照这篇文章：《tvm在linux环境下的安装与编译及vscode如何配置tvm的远程连接调试环境》进行配置，这里不在赘述。不过需要注意的是，在arm-ubuntu中并没有找到conda/build-environment.yaml文件中要求的llvmdev ==10.0.0这个版本，因此这里略作修改：

# conda/build-environment.yaml
# 这里将llvmdev更改为10.0.1版本
# 这样在编译时cmake会自动安装llvm
llvmdev ==10.0.1

  基于上述配置，再次修改build/config.cmake文件：

set(USE_LLVM ON)	# line 136
set(USE_ARM_COMPUTE_LIB ON)	# line 236
set(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR "/home/liyanpeng/arm64v8_work/acl") # line 237

  然后进行编译：

cd build
cmake ..
make -j6

  俩小时过去了。。。

  编译成功后的信息如下：

5.2 关于ubuntu 18.04升级glibc掉入坑里这件事

  遗憾的是，在pytorch 1.7.1版本及以前的版本中，官方并没有提供arm版本的pytorch，这里有两种解决方式：
  (1) 从社区下载非官方版本的pytorch-aarch64，比如：KumaTea
  (2) 官方从pytorch 1.8.0版本开始提供了arm版本的pytorch，可以选择更高版本的pytorch，不过还是要说一下，tvm官方目前[文章发布时]支持pytorch 1.7和pytorch 1.4两个大版本，其他版本可能不稳定。
  这里选择了非官方版本的pytorch-aarch64：

pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 -f https://torch.kmtea.eu/whl/stable-cn.html

  在查看pytorch版本时报错：ImportError: /lib/aarch64-linux-gnu/libc.so.6: version "GLIBC_2.28" not found[最初用的是ubuntu:18.04这个版本，换成ubuntu:20.04这个版本不会报错，可直接看5.3小节]

  查看系统当前glibc版本

ldd --version
# or
strings /lib/aarch64-linux-gnu/libm.so.6 | grep GLIBC_

# Ubuntu 18.04: 2.27
# Ubuntu 20.04: 2.31

  好了，可以打住了，建议掉头，直接跳到5.3小节，不然等会儿掉坑里可能出不来！！！

  解决方法可以参考这篇博客：

# 安装依赖
apt-get install gawk
apt-get install bison

apt-get install wget

# 下载、解压并配置
wget http://ftp.gnu.org/gnu/libc/glibc-2.28.tar.gz
tar -zxvf glibc-2.28.tar.gz
cd glibc-2.28
mkdir build
cd build
../configure --prefix=/usr/local --disable-sanity-checks

# 安装
make -j6
make install

  安装过程中的一些日志信息如下：

  没有任何报错信息，说明安装成功。

# 查看原始的软连接
ll /lib/aarch64-linux-gnu/libc.so.6

  按网上一些教程出现Segmentation fault错误，导致常用的ls，cp，clear这些指令无法使用，解决方法：

# export LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so:/lib/aarch64-linux-gnu/ld-2.27.so
unset LD_PRELOAD
# 取消软连接
LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so unlink /lib/aarch64-linux-gnu/libc.so.6
# 重新恢复
LD_PRELOAD=/lib/aarch64-linux-gnu/libc-2.27.so ln -s /lib/aarch64-linux-gnu/libc-2.27.so /lib/aarch64-linux-gnu/libc.so.6

  建立软连接：

# 复制 libc
cp /usr/local/lib/libc-2.28.so /lib/aarch64-linux-gnu/
cp /usr/local/lib/ld-2.28.so /lib/aarch64-linux-gnu/

cd /lib/aarch64-linux-gnu/
# ll ld-linux-aarch64.so.1
# ll libc.so.6

ln -sf /lib/aarch64-linux-gnu/libc-2.28.so /lib/aarch64-linux-gnu/libm.so.6
# 无效, 仍然是2.27版本

  对比了aarch64-linux-gnu目录与glibc-2.28的安装目录，发现好多库名一样，只是版本号不一样，是不是都要进行替换？？？

5.3 验证安装是否成功

  到了这一步，如果ubuntu 18.04的glibc没有升级成功，那就使用ubuntu 20.04吧，以下是在arm-ubuntu 20.04中进行的，查看pytorch版本：

  pytorch模型验证：

# from_pytorch.py

import tvm
from tvm import relay
from tvm.contrib.download import download_testdata
import numpy as np

import torch
import torchvision

######################################################################
# Load a pretrained PyTorch model
pth_file = 'resnet18-f37072fd.pth'
model = torchvision.models.resnet18()
ckpt = torch.load(pth_file)
model.load_state_dict(ckpt)
model = model.eval()

# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()

######################################################################
# Load a test image
from PIL import Image

img_path = 'cat.png'
img = Image.open(img_path).resize((224, 224))

# Preprocess the image and convert to tensor
from torchvision import transforms

my_preprocess = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)

######################################################################
# Import the graph to Relay
input_name = "input0"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

######################################################################
# Relay Build
target = tvm.target.arm_cpu()
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)

######################################################################
# Execute the portable graph on TVM
from tvm.contrib import graph_executor

dtype = "float32"
m = graph_executor.GraphModule(lib["default"](dev))
# Set inputs
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)

#####################################################################
# Look up synset name
synset_path = 'imagenet_synsets.txt'
with open(synset_path) as f:
    synsets = f.readlines()

synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = {
    
    spl[0]: " ".join(spl[1:]) for spl in splits}

class_path = 'imagenet_classes.txt'
with open(class_path) as f:
    class_id_to_key = f.readlines()

class_id_to_key = [x.strip() for x in class_id_to_key]

# Get top-1 result for TVM
top1_tvm = np.argmax(tvm_output.numpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]

# Convert input to PyTorch variable and get PyTorch result for comparison
with torch.no_grad():
    torch_img = torch.from_numpy(img)
    output = model(torch_img)

    # Get top-1 result for PyTorch
    top1_torch = np.argmax(output.numpy())
    torch_class_key = class_id_to_key[top1_torch]

print("Relay top-1 id: {}, class name: {}".format(top1_tvm, key_to_classname[tvm_class_key]))
print("Torch top-1 id: {}, class name: {}".format(top1_torch, key_to_classname[torch_class_key]))

  验证结果如下【速度不是有一点慢】：

  autotvm也是ok的：

  但感觉有些问题，在x86上task只有13个[resnet18]，而在arm上确是26个，经查验这26个task即包含arm的，还包含x86，就很奇怪，这个问题还待解决：

结束语

  本篇文章是对在x86的docker中搭建tvm的arm环境的一次尝试，由于服务器上使用的是ubuntu 18.04，本想着在本地模拟一个同服务器相同的版本，结果遇到了ubuntu 18.04升级glibc这个坑，暂时还没有填上。不过，成功在x86的docker中搭建了tvm的arm环境，并能进行编译和运行。
  tvm仍在持续学习中，欢迎评论区交流哦！！！