TVM学习

【TVM系列二】TVM介绍

2022-07-02  本文已影响0人  AIPlayer

一、TVM的工作流程

TVM主要由两个部分组成：

（1）TVM编译器：负责编译和优化模型

（2）TVM runtime：提供目标设备上运行模型的API

1、整体流程
image.png

如图所示，TVM的工作流程包括4个主要部分：

2、关键数据结构

IRModule (intermediate representation module)是贯穿整个TVM的数据结构，重要性不言而喻。它是一系列Function的集合，用于表述一个神经网络模型，目前TVM支持两种主要的变体函数：

在整个编译过程中，一个relay function可能会被优化为多个tir::PrimFunc。

3、Transformations

transformation的作用有两个：

（1）优化（optimization）：将程序转换为等效的、或者更优化的版本；
（2）底层表示（lowering）：将程序转换为更接近目标设备的低层级表示。

relay/transform 包含一组优化模型的passes。优化包括constant folding和dead-code消除，以及针对张量计算的优化，如layout转换和scaling factor folding。

在relay优化的pipeline的最后，会运行一个 FuseOps的pass，将一个完整的Function（对应一个端到端的模型如 MobileNet）分解为多个子Funcions（例如 conv2d-relu）段。这样做的好处是将问题分成了两个子问题：

tir/transform 包含 TIR 层级functions的转换passes。例如，将multi-dimensional access flatten到一维访问，将内在函数扩展为特定于目标的函数，以及修饰函数入口以满足运行时调用约束。除此之外，也有优化passes，如access index简化和dead-code消除。

4、搜索空间和基于机器学习的转换

前面描述的转换都是基于规则和确定的，而TVM的设计目标之一是支持对于不同的硬件平台都可以进行高性能的代码优化。因此，需要对尽可能多的优化进行选择，每个优化又需要选择最优的参数，从这个角度来看，这个的工作量无疑是巨大的。TVM采用了基于空间搜索和机器学习的方法来解决这个优化选择与调参的问题。

顾名思义，空间搜索需要在特定的空间，所以首先需要定义一系列的转换操作，比如循环转换、内联、矢量化等。这些操作称为调度原语（scheduling primitives）。调度原语的集合定义了可以对程序进行优化的搜索空间，然后TVM搜索不同的调度顺序以挑选最佳调度组合。这个搜索过程通常由机器学习算法完成，TVM使用的是xgboost算法。在搜索完成后，记录下每个算子最优的调度顺序，然后编译器就可以将此调度序列应用到程序中。TVM使用基于搜索的优化方法来处理初始 tir function生成问题。这部分模块称为 AutoTVM(auto_scheduler)。

5、目标代码转化

目标代码转换阶段主要是将 IRModule 转换为可以相应的目标设备上运行的格式：

代码生成阶段需要尽可能地轻量化，所以绝大多数的转换和降层级都应该放在目标代码转换之前执行。

二、逻辑架构组件

image.png

上图显示了TVM中的主要逻辑组件：

三、运行TVM实例

1、交叉编译runtime

想要在目标设备上运行模型的前提是交叉编译模型和runtime库。以Raspberry Pi为例，首先需要在主机安装Raspberry Pi的编译工具链：

sudo apt-get update
sudo apt-get install gcc-aarch64-linux-gnu g++-aarch64-linux-gnu
sudo apt-get install gcc-multilib-arm-linux-gnueabihf g++-multilib-arm-linux-gnueabihf

然后交叉编译TVM runtime库：

cmake .. \
    -DCMAKE_SYSTEM_NAME=Linux \
    -DCMAKE_SYSTEM_VERSION=1 \
    -DCMAKE_C_COMPILER=/usr/bin/aarch64-linux-gnu-gcc \
    -DCMAKE_CXX_COMPILER=/usr/bin/aarch64-linux-gnu-g++ \
    -DCMAKE_FIND_ROOT_PATH=/usr/aarch64-linux-gnu \
    -DCMAKE_FIND_ROOT_PATH_MODE_PROGRAM=NEVER \
    -DCMAKE_FIND_ROOT_PATH_MODE_LIBRARY=ONLY \
    -DMACHINE_NAME=aarch64-linux-gnu

make -j2 runtime

编译完成后使用file命令查看编译出来的runtime库是否OK：

image.png
2、编译模型

在主机上构造一个简单的kernel，并在主机上编译，示例代码如下：

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)
# 编译并保存结果：local_demo为True表示编译target为主机端运行，否则为raspbarry pi
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") # 为目标设备生成代码
print(func)
path = "./tvm_test_lib.tar"
func.export_library(path)

运行代码后会得到tvm_test_lib.tar的编译结果，func的打印输出为：

Module(llvm, 56334d7e8738)

它是一个 tvm.runtime.PackedFunc 类型，TVM使用Function开作为前后端的黏合，一个编译后的module返回Function，TVM后端同样也以Functions的方式注册和暴露API。

3、运行模型

将它用rpc的方式运行在设备上，需要将编译的lib上传到设备，然后使用设备端的编译器重新链接之后，func就是一个设备端的模型对象了。

if local_demo:
remote = rpc.LocalSession()
else:
    host = "192.168.1.111"
    port = 9090
    remote = rpc.connect(host, port)

remote.upload(path)
func = remote.load_module("tvm_test_lib.tar")
print(func)
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_f = func.time_evaluator(func.entry_name, dev, number=10)
cost = time_f(a, b).mean
print("%g secs/op" % cost)

此时的func打印输出是：

Module(rpc, 56334d6d7148)

四、总结

本文介绍了TVM的工作流程和内部的逻辑框架组件，通过运行TVM的一个实例了解和熟悉TVM的Python API使用。

上一篇 下一篇

猜你喜欢

热点阅读