深入浅出 tvm – (5) IR 公共的一些核心概念

在深入 OptimizeImpl 阶段,也就是 Pass 优化之前,我们先了解一下 Relay 阶段的一些基本概念
OptimizeImpl 阶段最主要的工作就是图优化,图其实就是一种高级的表示,tvm 和图相关的概念好好几个:
  1. relay ir:这是 tvm 最 high level 的表示,另外 relay ir -> tir 的过程中,又会依赖 topi 和 te 这2种特定抽象的中间表示
    1. topi:TVM Operator Inventory,TOPI 提供了比 tir 具有更高抽象的 numpy 风格的,通用操作和调度
    1. te:Tensor Expression,张量表达式
  1. tir:最接近目标代码的中间表示
  2. ir:relay ir 和 tir 的一些公共基础结构,和上述2种ir不一样,并不是一个完整独立的抽象
本章我们先来了解下 IR 这个 relay ir 和 tir 最公共的基础设施,后续会依次介绍 relay ir、tir、topi、te
代码目录:
  • 代码:src/ir
  • 头文件:include/tvm/ir
编程语言最基本的核心概念就3个:类型、运算符、表达式,在 IR 这里分别对应 Type, OP, Expr

1.1 Type

Type 相关的定义都在 include/tvm/ir/type.h ,Type 包括基础的整型/浮点型等,也包括函数类型等相对复杂的类型。
这里我们介绍2种基本的类型:
  1. PrimType:最原始的 Type,可以直接映射到 low-level IR 的基本数据类型
  2. FuncType:函数类型
PrimType 可以在这上面做一些 Low-level 优化
定义如下:
class PrimTypeNode : public TypeNode {
 public:
  /*!
   * \brief The corresponding dtype field.
   */
  runtime::DataType dtype;
}
可以看到 PrimType 就一个数据成员,runtime::DataType,这个是 runtime 最底层的概念,代码在 include/tvm/runtime/data_type.h
/*!
 * \brief Runtime primitive data type.
 *
 *  This class is a thin wrapper of DLDataType.
 *  We also make use of DataType in compiler to store quick hint
 */
class DataType {
 public:
  /*!
   * \brief Type code for the DataType.
   *
   * DLPack consistency:
   * 1) kInt is consistent with kDLInt
   * 2) kUInt is consistent with kDLUInt
   * 3) kFloat is consistent with kDLFloat
   */
  enum TypeCode {
    kInt = kDLInt,
    kUInt = kDLUInt,
    kFloat = kDLFloat,
    kHandle = TVMArgTypeCode::kTVMOpaqueHandle,
    kBFloat = kDLBfloat,
    kCustomBegin = 129
  };
  /*! \brief default constructor */
  DataType() { data_ = DataType::Void(); }
  /*!
   * \brief Constructor
   * \param dtype The DLDataType
   */
  explicit DataType(DLDataType dtype) : data_(dtype) {}


深入浅出 tvm – (4) Relay 计算图

在深入了解 IR 以及 relay.IR 之前,我们先对计算图在 relay 里的表示有一个直观的认识
幸运的是,relay 提供一个可视化组件 Relay Visualizer,帮我们了解其计算图的内部结构:https://tvm.apache.org/docs/how_to/work_with_relay/using_relay_viz.html
获得计算图有2种方式:
  1. 加载开源模型,tvm提供一些便捷的函数
  2. 自定义手写一个模型
我们看下

1. 开源模型

我们可以通过 relay.testing 模块获取一些常见的模型,具体可以看 python/tvm/relay/testing 模块,目前支持的模型有:
  1. resnet
  2. resnet_3d
  3. mobilenet
  4. mlp
  5. lstm
  6. synthetic
比如,我们想看一下 resnet18 模型的图结果,可以这么看
from tvm import relay
from tvm.relay import testing
import tvm
from tvm.contrib import relay_viz

# Resnet18 model
resnet18_mod, resnet18_params = relay.testing.resnet.get_workload(num_layers=18)

viz = relay_viz.RelayVisualizer(
    resnet18_mod,
    relay_param=resnet18_params,
    plotter=relay_viz.DotPlotter(),
    parser=relay_viz.DotVizParser())
viz.render('1')
viz.render 会在当前目录生成 1.pdf 的文件,打开即可得到完整的计算图,如下:
0
这个计算图非常大,这里只展示部分,从这个图里面我们可以看到,整个图基本上只有2类节点:
  1. Var:本地变量
  2. Call:算子调用
除了 Var, Call 之外,Relay 的计算图还包含 Function, GlobalVar, Tuple 等其他节点类型,后面我们将深入展开描述


深入浅出 tvm – (3) 架构 & 设计

本文将从一个最简单的图像分类模型,了解基本的推理过程,并以此来了解一下 tvm 编译器的基本工作流程,以及对应核心模块的职责
0

1. 模型的定义

由于tvm支持onnx模型格式,所以我们从 resnet50-v2-7.onnx 开始
模型的输出:https://s3.amazonaws.com/onnx-model-zoo/synset.txt,其实就是1000个图像分类
当我们执行模型推理时,输入一个图像,输出是一个 shape=(1,1000) 的数组,表示这个图片对应1000个分类的概率,其中概率最大的那个分类,就是预测的结果

2. 从一个最基本的推理开始

推理的过程就是把模型,编译成目标设备上可运行的代码,根据输入数据,返回预测结果
如下,说下几个关键的地方:
  1. 模型转换:relay.frontend.from_onnx(),把模型从 onnx 格式转换成 relay IRModule 格式
  2. 编译:relay.build(),完成从IRModule到目标设备代码的编译优化,内部包含了 pass 优化、schedule优化
  3. 推理:输入 module.set_input(),运行 module.run(),获取结果 module.get_output()
其中最核心的就是2这几个步骤,基本包括了tvm最核心的逻辑
# https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx
onnx_model = onnx.load('./resnet50-v2-7.onnx')
img_data = get_test_img()
input_name = "data"
shape_dict = {input_name: img_data.shape}

mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target="llvm", params=params)

module = graph_executor.GraphModule(lib["default"](tvm.device("llvm", 0)))

module.set_input(input_name, tvm.nd.array(img_data.astype('float32')))
module.run()
"https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
out = module.get_output(0, tvm.nd.empty((1, 1000))).numpy()
我们再来看下 relay.build 具体干了啥:
  1. python/tvm/relay/build_module.py: build()
    1. python/tvm/relay/build_module.py:BuildModule::build()
      1. src/relay/backend/build_module.cc: BuildRelay()
BuildRelay 干的事情就是:Compile a Relay IR module to runtime module


深入浅出 tvm – (2) 搭建开发环境

本文面向 tvm 开发者,粗略介绍一下 tvm 的安装过程

1. 开发环境

系统:ubuntu 20.04,由于中文生态兼容的比较好,比如输入法,有道云笔记,所以这个系统本身也是我的主力开发系统
tvm 版本:v0.10.0,如果你的不是这个版本,记得 checkout 一下,否则下面的依赖可能不全,对于学习 tvm 来说,v0.10.0 版本其实够了

2. 依赖安装

以我的 dockerfile 为例,可以参考
另外,大家也可以跳过这一步,直接用我弄好的镜像:docker pull pipul/tvm_ci_cpu
FROM ubuntu:20.04

ENV TZ=Europe/Kiev
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

COPY ./apt /etc/apt
COPY ./etc/sudoers.d/90-nopasswd-sudo /etc/sudoers.d/90-nopasswd-sudo
COPY ./.bashrc /etc/bashrc

RUN apt-get update
RUN apt-get install -y vim sudo git g++ pip llvm
RUN apt-get install -y python3 python3-dev python3-setuptools gcc libtinfo-dev zlib1g-dev build-essential cmake libedit-dev libxml2-dev
RUN pip install onnx numpy attrs decorator Pillow psutil scipy relpath typing-extensions tornado 'xgboost>=1.1.0' cloudpickle
有几个不是必须的:
  1. COPY ./apt 主要是替换了163的源,不然软件安装太慢
  2. /etc/sudoers.d/90-nopasswd-sudo 是 tvm 编译的时候必须依赖的,但是 docker 镜像 ubuntu:20.04 是默认没有的,所以需要搞下
  3. COPY ./.bashrc 可以忽略
假设构建好的镜像为:${tvm-docker-image}
一旦镜像构建好了之后,后续编译和开发 tvm,都在镜像内完成


深入浅出 tvm – (1) 简介

深度学习分为2个过程:
  1. 训练:从数据得到模型
  2. 推理:从模型得到答案
为解决训练而设计的系统叫训练框架:比如 paddlepaddle,tensorflow,pytorch
为解决推理而设计的系统叫推理引擎:比如 paddle inference,tensorflow,pytorch-trt 等等
训练得到的模型,其实就是一个计算图,这个计算图接收输入,通过一系列的运算,得到一个结果,后面这个过程就叫推理

1. 编译器在什么位置?

如果只从系统的输入和输出看的话:编译器 ≈ 推理引擎
因为现在几乎所有的AI编译器,它的输入都是模型(计算图)+数据,输出就是一个推理结果,这个和推理引擎所干的事情几乎一模一样
以 tvm 为例,如下:
0
其中 nnvm + graph optimizations + tvm + tvm primitives 就是 TVM 干的事情
tvm 的输入就是 CoreML 或者 ONNX 模型(一种计算图的格式)