深入浅出 tvm – (7) Relay IR 之 ADT

ADT 的全称是 Algebraic data type
ADT 是函数式编程的一个重要特征,也是 tvm relay 的一个高级特性。但是目前 ADT 看起来还是 tvm 的一个内部能力,还不能像我们构建计算图一样直接用文本描述出来,因此官方文档里面的示例都是只用来参考的
由于目前还不能通过调试的方式,深入的了解 ADT,所以本文只能算是对官文文档的一个解读,方便大家对 ADT 有一个初步的概念
参考:
0

1. ADT 的基本概念

1.1. 定义和匹配

简单来说,一个ADT定义含有多个构造函数,每个构造函数有不同的参数类型。每个ADT实例构造出来后只是简单地包含其构造时的所有参数值
重点:最后一句。这个和 c++ 的结构体还不太一样。在c++里面,可以允许有多个构造函数,但是结构体本身是不变的,成员可以被默认初始化。但是在这里,ADT 你可以理解为就是一个 tagged unions
如下:
Numbers 是一个 ADT,具有3个构造函数
# Defines an ADT named "Numbers"
data Numbers {
  Empty : () -> Numbers
  Single : (Tensor[(), int32]) -> Numbers
  Pair : (Tensor[(), int32], Tensor[(), int32]) -> Numbers
}
# A Numbers value can be produced using an Empty, Single, or Pair
# constructor, each with a signature given above
由于 ADT 的值,在析构之前是不确定的,因此,当一个函数接收 ADT 实例作为参数时,我们必须根据其构造函数的类型,来决定下一步做什么
这就有了 match 语法,如下:
def @sum(%n : Numbers[]) -> Tensor[(), int32] {
   # The match expression branches on the constructor that was
   # used to produce %n. The variables in each case are bound
   # if the constructor matches that used for %n
   match(%n) {
     case Empty() { 0 }
     case Single(x) { x }
     case Pair(x, y) { x + y }
   }
}

@sum(Empty())    # evaluates to 0
@sum(Single(3))  # evaluates to 3
@sum(Pair(5, 6)) # evaluates to 11
由于 ADT 是通过名字标识的,意味着2个相同具有相同构造函数的ADT,在类型系统看来仍然是不同的
比如下面,调用 @sum(Empty2()) 就会报错
# structurally identical constructors to Numbers
data Numbers2 {
  Empty2 : () -> Numbers2
  Single2 : (Tensor[(), int32]) -> Numbers2
  Pair2 : (Tensor[(), int32], Tensor[(), int32]) -> Numbers2
}

# the below results in a type error because Numbers2
# is a distinct type from Numbers
# fn() { @sum(Empty2()) }


深入浅出 tvm – (6) Relay IR:一种 high-level 的中间表示

前面一章我们已经了解了 relay ir 和 tir 中公共的 ir 概念,现在我们继续深入一下 relay ir
relay ir 是 tvm 最高层次的 IR,也是最接近前端的一种计算图表示方式
相比传统的基于 Graph 的 High Level IR, Relay 引入了不少函数式编程的概念,带来了强大的表达能力。
但是相比官方文档,这个写的更好:https://zhuanlan.zhihu.com/p/446976730

1. Type

include/tvm/relay/type.h
relay 并没有定义任何新的 Type 类型,全部直接服用 tvm/ir 里的定义,因此这里不展开描述
// namespace update for backward compact
// will be removed later.
using AnyNode = tvm::tir::AnyNode;
using Any = tvm::tir::Any;
using Kind = TypeKind;
using Type = tvm::Type;
using TypeNode = tvm::TypeNode;
using TypeVar = tvm::TypeVar;
using TypeVarNode = tvm::TypeVarNode;
using GlobalTypeVar = tvm::GlobalTypeVar;
using GlobalTypeVarNode = tvm::GlobalTypeVarNode;
using TupleType = tvm::TupleType;
using TupleTypeNode = tvm::TupleTypeNode;
using TypeConstraint = tvm::TypeConstraint;
using TypeConstraintNode = tvm::TypeConstraintNode;
using FuncType = tvm::FuncType;
using FuncTypeNode = tvm::FuncTypeNode;
using IncompleteType = tvm::IncompleteType;
using IncompleteTypeNode = tvm::IncompleteTypeNode;
using RelayRefType = tvm::RelayRefType;
using RelayRefTypeNode = tvm::RelayRefTypeNode;
using TensorType = tvm::TensorType;
using TensorTypeNode = tvm::TensorTypeNode;
using TypeCall = tvm::TypeCall;
using TypeCallNode = tvm::TypeCallNode;
using TypeRelation = tvm::TypeRelation;
using TypeRelationNode = tvm::TypeRelationNode;
using TypeRelationFn = tvm::TypeRelationFn;
using TypeReporter = tvm::TypeReporter;
using TypeReporterNode = tvm::TypeReporterNode;

2. Expr

relay 里面的所有 Expr 都是继承自 ir 里面的 RelayExpr 类
namespace tvm {
namespace relay {

using Expr = tvm::RelayExpr;
relay 定义的 Expr 列表如下:
  1. Constant:常量
  2. Tuple:数组,由N个Expr构成
  3. Var:局部变量,计算图最常见的结构之一
  4. Call:算子调用,计算图最常见的结构之一
  5. Let:Let binding
  6. If:条件表达式
  7. TupleGetItem
  8. RefCreate
  9. RefRead
  10. RefWrite
  11. TempExpr
比如 Tuple
class Tuple : public Expr {
 public:
  /*!
   * \brief The constructor
   * \param fields The fields of a tuple.
   * \param span The source span of the expression.
   */
  TVM_DLL explicit Tuple(tvm::Array fields, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Tuple, RelayExpr, TupleNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(TupleNode);
};

其他的比如 Constant/TupleGetItem 这样比较简单的 Expr, 就不介绍了,有兴趣的同学可以阅读 include/tvm/relay/expr.h
再来看下 Var 的定义
class VarNode : public ExprNode {
 public:
  /*!
   * \brief The unique identifier of the Var.
   *
   * vid will be preserved for the same Var during type inference
   * and other rewritings, while the VarNode might be recreated
   * to attach additional information.
   * This property can be used to keep track of parameter Var
   * information across passes.
   */
  Id vid;
  /*!
   * \brief type annotaion of the variable.
   * This field records user provided type annotation of the Var.
   * This field is optional and can be None.
   */
  Type type_annotation;
  // ...
}
Id 可以简单理解为一个字符串,是一个全局唯一的标识。你可以改下代码,增加一些调试,打印出来是这样的:
fc1_weight
fc1_bias
stage1_unit1_bn1_gamma
stage1_unit1_bn1_beta
另外这里有个地方没太理解,Var的值是存在哪里的?


深入浅出 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 模型(一种计算图的格式)


Efficient Memory Disaggregation with Infiniswap

1. 概要

infiniswap 可以实现把一个机器的内存 swap 到另外多个机器上(1对多)
和其他实现相比,infiniswap 可以把开销做到很小,并且故障容错能力更强
  1. A transparent remote paging model for virtual machines @2008
  2. Swapping to remote memory over Infiniband: An approach using a high performance network block device @2005
  3. 还有一堆20+年前的论文(参考意义不大)

2. 架构设计

infiniswap 的组件其实就2个,非常简单:
  1. infiniswap-bd:一个虚拟的 swap 设备,其后端是远程内存
  2. infiniswap-daemon:管理本机对外提供的 remote memory(包括内存到虚拟地址的映射,等等)
这2个组件每个机器都会部署,类似于 daemonset
而且有个比较有意思的地方是,infiniswap 是去中心化的,它并没有一个中心式的管理组件(你看zombieland,中心式的管控组件就挺多的)
如下:
0


Zombieland: Practical and Energy-efficientmemory disaggregation in a datacenter

Zombieland: Practical and Energy-efficientmemory disaggregation in a datacenter

1. 概要

zombieland 从硬件视角,实现了CPU&内存的池化
和传统的专门做内存池化的系统不太一样,zombieland 做到了把cpu和内存彻底的分离,甚至连供电都是分开的,它可以做到把cpu的电源关掉的情况下仍然可以让内存提供远程访问服务
(有点牛逼,这不得把主板设计都改了?)
通过这个设计,zombieland 实现了 86% 的能效提升

2. 设计实现

zombieland 的思路是实现服务器的 ”半关机“,增加了一个新的 ACPI 状态,叫 Sz,处于这种状态的服务器叫 zombie
这种 Sz 服务器,它的内存是可以被远程访问的,但是功耗只比 S3 状态增加了 15%,而且这种设计对硬件的改动很小
然后 zombieland 给 hypervisor 实现了2种远程内存扩展:RAM ext 和 Explicit swap device,允许 hypervisor 把虚机的冷内存卸载到远程内存上,或者用户可以显式的把内存swap到远端

2.1. 理论收益

和其他方案对比:
1)Server-centric architecture
0
能效:Total Energy Consumed = 2.1 × Emax.


AIFM – High-Performance, Application-Integrated Far Memory

AIFM 提出一种在用户态完成 swap 的方案,相比传统的内核态(虚拟内存)的方案,实现了:
  1. 更低的延迟:相比 fastswap,有6倍吞吐的提升
  2. 更高的灵活性:用户可以显式指定一块内存是local还是far(指向 remote memory)
0

1. 一些背景

1) OS swapping and far memory
之前对 swap 以及 far memory 相关的研究,大多都是在内核态的,由于内核态的 swap 都是 4k 粒度的,对于一些随机大量的小object场景,会带来很多性能问题
另一方面,如果你要依赖一些应用的信息(比如访存的行为、访存策略、内存热点)去做性能优化,这也是做不到的(madvise有一些近似的功能,但是不够)
为此,AIFM 提出了一种object级别的,可以访问远端内存的,类似于c++弱指针的一种新型指针
2)Disaggregated and distributed shared memory
当前大多数 Disaggregated memory 相关的实现,都要求数据中心具备高速互联的能力(但是现在很多数据中心其实还不具备这个能力),AFIM 提供了一种软件级的解决方案
另外一个问题是,传统的 far memory 系统里面,一份数据需要被多个 host 共享,这就要解决很复杂的数据一致性问题(同时带来了很大的性能开销)
AIFM 简化了这个问题, far memory 中的数据只属于一个 host