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


nr_dying_descendants 引发的在线业务性能抖动

最近发现一则比较奇怪的业务性能抖动问题

简单来说,就是离线作业频繁的触发oom,导致在线业务指标陡增。理论上来说,离线作业触发自身硬限,不应该导致其他容器出问题才是

先来说一下我们的 memory cgroup 层级结构,可以简单理解为如下:

/cgroups/v2/kubepods/

/cgroups/v2/kubepods/online/a -> a为在线服务

/cgroups/v2/kubepods/offline/job1 -> job1 为离线服务

通常来说,为了解决混部的内存争抢问题,我们会对整个 offline 大框设置一个内存上限,防止离线作业使用内存太多,导致在线服务出现内存压力。这看起来似乎没什么问题。

不过最近一个问题,发现离线 job1 频繁 oom,引发了在线的性能抖动

最后定位的结论如下:

  1. /cgroups/v2/kubepods/offline/cgroup.stat 下看到有大量的 nr_dying_descendants
  2. offline 大框频繁出发 oom
  3. oom 的回收过程有个效率不太好的地方,导致 oom 回收效率底下 -> 导致内核各种锁争抢严重 -> 进程sys高 -> cpu排队严重 -> 在线业务性能抖动

根因分析:nr_dying_descendants 的原因是由于容器被销毁时,容器内还有未完全回收干净的 page cache,所以内核不会立即释放这个 cgroup(用户态操作系统从 /cgroups 文件系统下已经看不到了,但是内核态的数据结构还在)

cat /cgroups/v2/kubepods/online/cgroup.stat

nr_descendants 10

nr_dying_descendants 4172

这就有一个问题,如果有大量的dying cgroups,内核在oom处理过程中:

  1. 如果是cgroup oom,会优先尝试从 dying cgroups 回收内存,但是最多只回收5个 -> 这个地方有效率问题
  2. 如果是整机回收,不处理dying cgroups内存

 


Concurrency management in BPF

eBPF 对并发编程的支持是非常弱的

有一些讨论:

  1. https://lists.iovisor.org/g/iovisor-dev/topic/bpf_concurrency/74407447
  2. https://lwn.net/Articles/779120/

但是你甚至在 eBPF 的官方文档 https://github.com/iovisor/bcc/blob/master/docs/reference_guide.md 里面几乎找不到任何关于并发的描述,除了 map.atomic_increment() 接口

但是这个借口也很奇怪。为了突出这个接口是原子的,线程安全的,难道其他接口就不是了吗?但好像又不是,因为 eBPF 里有大量的示例,并没有对 map 做特殊的并发控制,所以理论上可以认为,其他大部分 map 操作接口,也是线程安全的

线程安全验证

测试代码如下:

// calculate latency
int trace_run(struct pt_regs *ctx, struct task_struct *prev)
{
    u32 pid = bpf_get_current_pid_tgid();
    u64 *tsp;
    if (pid <= 1) {
        return 0;
    }
    
    pid = 10000; 
    u64 ts = bpf_ktime_get_ns();
    start.update(&pid, &ts);
   
#define __OP__ == 
    tsp = start.lookup(&pid); if (tsp __OP__ 0) { dist.increment(bpf_log2l(1)); }
    tsp = start.lookup(&pid); if (tsp __OP__ 0) { dist.increment(bpf_log2l(1)); }
    /** 重复N次 **/
    tsp = start.lookup(&pid); if (tsp __OP__ 0) { dist.increment(bpf_log2l(1)); }
    tsp = start.lookup(&pid); if (tsp __OP__ 0) { dist.increment(bpf_log2l(1)); }
    
    start.delete(&pid);
    return 0;
}

这段代码,修改自 bbc 工具集自带的 runqlat.py

这段代码实际上是在构造一个冲突场景,来验证 trace_run 到底是并行执行还是串行执行的。代码的逻辑是在一个超高频触发的钩子(调度器进行上下文切换)函数里,先在 map 里插入一个固定key(pid)的元素,然后查询 N 次,每次查询如果 key 不存在,计数器就 +1,最后把 key 从 map 里面删除掉,如果:

  1. trance_run 是串行的,那么技术器永远都是 0,因为在一个串行执行的逻辑里,插入元素后,不管查询多少次,元素必然是存在的
  2. 如果 trace_run 是并行的,那么就存在一个条件,一定有一个 cpu,在另外一个 cpu 查询 key 是否存在的时候,执行到了函数结束的地方,把 key 从 map 里面删掉了

事实证明,trace_run 是并行执行的,并且中间 lookup 的次数越多,冲突的概率就越大


cgroup 内存管理之 mlock 内存

我们在使用容器的过程中,可能会遇到一个问题,容器利用率很低,但是经常发生 oom,这是什么情况?

很可能是业务使用了一些不可回收的 page cache,其中最主要的应该就是 mlock

1. mlock 的背景

mlock 的作用就是防止页面被换出到 swap 分区,或者 page cache 被内核回收

由于线上服务器,基本默认都会关闭 swap 分区,所以这种情况暂不讨论,防止 page cache 被内核回首是主要作用

什么样的 page 需要 mlock?

1)通过 mmap 方式映射到内存中的 page cache,比如一些关键的索引数据,需要经常访问,那就得 mlock 住,否则 pgmajfault 带来的性能抖动是很大的

2)程序自身依赖的一些 so 动态链接库,由于内核的 page cache 回收算法,并不感知 page 具体是普通的文件 cache,还是 so 动态链接库,所以当容器内存不足时,内核通过一些粗略的回收算法回收 page cache,一旦把 so 的缓存页回收掉了,程序在调用相关函数时,会出现严重的性能抖动

因此,通过 mlock,显式的把一些关键的不希望被回收的 page cache 锁定起来,达到保证业务性能的目的

2. mlock 的内核实现

mlock 相关的系统调用可以参考:https://linux.die.net/man/2/mlock

2.1. unevictable 内存

cgroup 的 memory.stat 接口,包含一个容器内存使用的详细信息

cat /cgroups/v2/matrix/memory.stat 
anon 200723943424
file 158665015296
kernel_stack 0
slab 0
sock 229249024
shmem 4258582528
file_mapped 48234749952
file_dirty 8605696
file_writeback 0
inactive_anon 4258234368
active_anon 200570675200
inactive_file 73935654912
active_file 40680521728
unevictable 39942680576

其中 unevictable 表示不剔除的 page cache,mlock 属于 unevictable 的一部分,除此之外,unevictable 还有一些其他内存