源码研习 — TVM中的IR设计与技术实现

一、关键问题

TVM中的 IR 是什么，架构设计上分几层？

解答：TVM的整体结构图如下：

概念上，分为两层：上层为面向前端组网的Relay IR， 下层为面向LLVM的底层 IR。

但从设计实现上，底层通过 Object 元类实现统一的AST Node表示，借助一个 IRModule 贯穿上下层。个人理解，TVM的 IR 实现上其实只有一层，只是封装后在直观概念上分为上下层。

• IRModule里持有的是BaseFunction列表

• 上层 relay::Funtion继承自 BaseFunction

  官方解释：relay::Function对应于一个end2end的模型。可以简单理解为一个支持控制流、递归、以及复杂数据结构的计算图。

• 下层tir::PrimFunc也继承自BaseFunction

  官方解释：tir::PrimFunc包含了一些底层threading、vector/tensor的"指令"。通常为模型layers中的一个Op执行单元

• 在编译阶段，一个relay::Function可能会被lower成多个tir::PrimFunc。

TVM架构上主要包含了哪些核心模块和概念？

解答：如下是各个模块的交互图：

从编译流程上来看，涉及的核心数据结构有两个：

• IRModule：包含relay::Function和tir::PrimFunc
  • 此部分也是 Pass 策略的输入输出单元，即 IRModule → pass→ IRModule
  • 传送门：TVM 的 Relay IR设计
• runtime::Module：经过lowering之后，可执行期的基本单元，包含很多runtime::PackedFunc（可以理解为KernelFunc
  • 传送门：TVM Runtime System

编译时的Pass策略主要在IRModule数据结构层面进行，分为两方面：

• ruled-base：包括relay/transform和tir/transform
  • 前者多为上层“图”结构上Pass优化，比如常量折叠，fusion
  • 后者多为下层偏向编译器方面的Pass优化，比如prefetch注入，unrollLoop
• search-based：包括auto-schedule和auto-tvm

在前后端交互上，TVM将所有的核心数据结构都暴露到了Python前端，易用性和灵活性极强：

• 所有的核心对象都可以通过Python API直接构造和操作，比如IRModule
• 支持在前端自定义组合pass和transformation
• 通过TVM的API直接操作 IR，支持Python端写pass

IRMoule是什么样的？

• IRModule通过IRModuleNode管理元信息

• 核心成员：

  • Functions
    • 表示计算的函数单元，如Conv、log
    • Function内部有通过params、body关联Var
    • 概念上，对应与AST的Module
  • Global_var

  import tvm
  from tvm import relay
  import numpy as np
  # step 1: modeling
  m,n = 4, 2
  x = relay.var("x", shape=(m,n), dtype='float32')
  out = relay.nn.softmax(x)
  net = relay.Function([x], out)
  
  # step 2: build and lowering
  module = tvm.IRModule.from_expr(net)
  lib = relay.build(module, "llvm")
  
  # step 3: input tensor data
  ctx = tvm.cpu(0)
  x_t = tvm.nd.array(np.random.uniform(size=[m,n]).astype('float32'), ctx)
  runtime = tvm.contrib.graph_runtime.GraphModule(lib["default"](ctx))
  runtime.set_input("x", x_t)
  runtime.run()
  print(runtime.get_output(0))
  
  # print(net.body)
  '''
  fn (%x: Tensor[(4, 2), float32]) {
    nn.softmax(%x)
  }
  '''
  
  # print(module)
  '''
  def @main(%x: Tensor[(4, 2), float32]) {
    nn.softmax(%x)
  }
  '''
  

Relay的pass是如何实现和管理的？

解答：上面提到，概念上讲，TVM可以看做是分两层的：Relay层和tir层，通过IRModule来贯穿。在Pass优化上，TVM也进行了两层的设计：

• 上层基于“图”的优化

  这部分很类似Paddle的pass，主要通过对 AST 的分析，应用一些上层的pass策略，主要包括：

  • 常量折叠、DSE、Layout转换、scaling因子折叠
  • 最后会应用fuse pass。比如将一个MobileNet表示成很多conv2d-relu 的“段”
  • pass的定义见relay/transform

• 下层基于“target”的优化

  这部分pass主要涉及 lowering到target时需要采取的优化策略，比如如何生成高效执行conv2d-relu的代码。主要包括：

  • Prefetch语句注入、VectorizeLoop、UnrollLoop、RemoveNoOp
  • SkipAssert、ThreadSync、HoistIfThenElse等
  • 此部分 pass有的可以直接复用底层编译器的pass，如LLVM、CUDA C等编译器。因此TVM主要关注和ML相关、且底层编译器未考虑到的场景

TVM 的 pass是通过遍历AST，进行node修改来实现的（类似paddle的动转静），通过TVM_REGISTER_GLOBAL注册和暴露支持的pass。

对于开发者来讲，TVM是如何便捷地支持新增一个Pass的呢？

TVM官方给出了一个[常量折叠 Pass的文档](Adding a Compiler Pass to Relay)。由于 TVM 的 IR 比较像AST，因此pass的新增主要包括如下几个步骤：

• 需要一个 AST Traversers

  用于确定哪些node是需要修改。在常量折叠pass中，实现了ConstantChecker，通过map结构的memo_记录哪些node是常量node。这里只涉及两个node的函数重载：ConstantNode和TupleNode

• 需要一个Expression Mutators

  用于修改和替换满足条件的node。在常量折叠pass中，只有三种node涉及折叠：LetNode、TupleItemGetNode和CallNode，因此也需要重载这三个函数即可

TVM的pass设计思想和架构，可以更多的参考Pass Infrastructure文档介绍。整体上借鉴了很多LLVM的pass设计思想。目标很明确，旨在实现如下效果：

• 可以灵活地排布Optimization单元，支持用户随意地进行pass piplines定制
• 提供友好地pass budug体验
• 避免用户去手动处理pass之间的依赖
• 简化开发者新增pass的流程，支持在python端写pass

TVM Pass实现上，可以分为三大类：

• Module-Level Pass
  • 利用全局信息进行优化，可以删减Function，如 DSE pass
  • 核心pass函数是PackedFunc类型，因此支持python、C++去写pass
• Funtion-Level Pass
  • 对Module中的每个Function进行优化，只有局部信息
  • 不允许删减Function
  • 如公共子表达式替换、vectorization
• Sequential-Level Pass
  • 顺序执行一系列的pass

FusionPass的基本原理：

• 会先将IRModule转为Graph

TVM 中的 auto-tvm的角色是什么？

解答：上面我们介绍的TVM的pass都是rule-based的，意味着开发者在新增pass时，其实是只要匹配什么样的模式，然后替换成什么样的模式。

这导致两个问题：

• pass的数量会很受限
• pass都需要预定义后才能支持

auto-tvm会先定义一些粒度比较小的优化策略，TVM会启发式组合应用、评估这些策略带来的提升，最后使用最佳的组合策略，以实现auto。

Relay结构是执行期的结构么？

解答：Relay的解释器（Interpreter）可以执行relay的表达式，但不适合生产环境部署时使用。原因是：

• 解释器是通过遍历 AST 来执行程序，遍历过程是很低效的。

• 无法友好支持动态代码。比如动态schduling、动态Tensor shape、还有控制流。解释器提供了简单的实现方案，但无法高效地编译和优化

  静态的代码优点：graphs是固定的，方便大刀阔斧地进行优化，比如内存静态分配，最佳的内存复用等。

TVM 也使用了 graph runtime技术——提供了一种快速执行机制，但仅支持部分Relay的programs

因此，Relay引入了 Virtual Machine，旨在取得部署、执行Relay programs时，性能与灵活性之间的平衡。

从用户的角度，可以通过relay.crete_executor(kind, ctx, target)接口来创建不同的执行器：

• kind取值为：graph、vm、debug
• 统一实现了evalutae(expr, *args)接口

前置知识：VM

• 传统的VM主要操作部分scalar和大量低阶instructions
• 对于ML，主要是Tensor，以及部分的高阶instructions
  • 耗时集中在计算密集型Op的调用，如GEMM和Conv
• 设计的核心点是：指令集的选择、指令表示
  • op-code 和 data payload

TVM中的VM的指令集的设计：

• 偏向high-level的设计，尽量与Relay层的operation相呼应
  • AllocTenor、If、Goto
• 核心的三种object对象：
  • NDArray、ADT 和 Closure，分别用于表示Tensor、tuple/list、closure data。
• 栈（Stack）和状态（State）
  • 栈帧用于标记当前的函数调用
  • 每个函数的寄存器都是在连续空间上申请的
• dispatch loop
  • VM实现了switch 和 goto

TVM 的VM compiler设计：

• 作用：将 Relay的IR 编译成字节码序列，即 tvm::relay::Module → tvm::relay::vm::Executable→ tvm::relay::vm::Function→tvm::relay::vm::VirtualMachine

TVM 的 VM 对序列化和反序列化的支持：

• Graph Runtime方案中序列化的结果是：
  • 权重参数保存为 .weight文件
  • graph保存为 .json文件
  • 计算kernel保存为.so库
• VM 方案中序列化的结果为：
  • Relay的 object文件 .o文件
  • 计算kernel保存为.so库

TVM的Runtime模块是什么样的？

解答：先看一个用户侧使用的接口样例：

import tvm
# Example runtime execution program in python, with type annotated
mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")
arr: tvm.runtime.NDArray = tvm.nd.array([1, 2, 3], ctx=tvm.gpu(0))
fun: tvm.runtime.PackedFunc = mod["addone"]
fun(a)
print(a.asnumpy())

Runtime时期的三大核心概念：

• runtime.Module：封装编译DSO的核心单元，包含了很多PackedFunc，可以根据name获取函数
• runtime.PackedFunc：后端生成的函数，对应于DL中的KernelFunc
• runtime.NDArray：封装了执行期的Tensor概念

TVM的target过程做了什么事情？

解答：这个应该是比较明确的，类似很多开源的框架，TVM会将 IRModule emit 到后端编译器去in-memory地生成可执行代码。

个人理解，target的过程涉及到编译，这对框架要求很高，在大多数场景下，这个过程应该是超级轻量级的，速度应该越快越好。

通过本地编译安装和试用TVM，发现target的过程超级快，几乎瞬发返回可执行函数。

TVM中编译执行和预测部署是什么样的？

解答：首先需要进行网络的定义：

import tvm
import numpy as np

n = 12
A = te.placeholder((n,), name="A") # Tensor
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B") # Tensor
C = te.compute(A.shape, lambda *i: A(*i) - 1.0, name="C") # Tensor

s = te.create_scheduleC[B.op, C.op])  # schedule
add_func = tvm.build(s, [A, B, C], "llvm", name="add") # compile

# prepare data
ctx = tvm.cpu(0)
a_t = tvm.nd.array(np.random.uniform(size=nn).astype(A.type), ctx)
b_t = tvm.nd.array(np.zeros(nn, dtype=A.dtype), ctx)
c_t = tvm.nd.array(np.zeros(nn, dtype=A.dtype), ctx)
add_func(a_t, b_t, c_t)

对于预测部署，可以将计算逻辑编译为DSO：

from tvm.contrib import cc

# serialization
add_func.save('./add_kernel.o')
cc.create_shared('./for_infer.so', ['./add_kernel.o'])

# load for inference
m = tvm.runtime.load_module('./for_infer.so')
add_func = m['add']  # load add kernel func
add_func(a_t, b_t, c_t)  # infer

对于model的序列化和加载的例子：

# Resnet18 workload
resnet18_mod, resnet18_params = relay.testing.resnet.get_workload(num_layers=18)
# build
with relay.build_config(opt_level=3):
    _, resnet18_lib, _ = relay.build_module.build(resnet18_mod, "cuda", params=resnet18_params)

# export library
file_name = "./deploy.so"
resnet18_lib.export_library(file_name)

# load it back
loaded_lib = tvm.runtime.load_module(file_name)
#infer
data = np.random.uniform(-1, 1, size=input_shape(mod)).astype("float32")
ctx = tvm.gpu()
gmod = graph_runtime.GraphModule(loaded_lib["default"](ctx))
gmod.set_input("data", data)
gmod.run()
out = gmod.get_output(0).asnumpy()

TVM中对训练是如何支持的？

解答：TVM支持训练包括如下几个核心模块

1. 自动微分 auto-diff

   TVM中提供了grads = te.gradient(out, inputs)接口，实现反向梯度的自动求导。但目前仍然是只是一个实现性功能

TVM的动态shape是如何实现的？

解答：理解TVM的动态shape实现机制，首先我们先看下：从用户的角度，动态shape怎么使用。

import tvm
import numpy as np

# 组网
n, m = te.size_var("n"), te.size_var("m")
A = te.placeholder((n,m), name="A")
k = te.reduce_axis((0, m), "k")
B = te.compute((n,),lambda i:te.sum(A[i,k], axis=k), name="B")
# 编译
s = te.create_schedule(B.op)
net = tvm.build(s, [A, B, n, m])
# 执行
def run(n, m):
  ctx = tvm.cpu(0)
  a = tvm.nd.array(np.random.uniform(size=[n,m]).astype(A.dtype), ctx)
  b = tvm.nd.array(np.zeros((n,)).astype(A.dtype), ctx)
  return net(a, b, n, m)

run(4, 6)
run(10, 16)

TVM提供了便捷的debug机制，可以直接打印查看中间编译的函数代码：

print(str(tvm.lower(s, [A, B])))

"""
primfn(A_1: handle, B_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B} {
  for (i: int32, 0, n) {
    B_2[(i*stride)] = 0f32
    for (k: int32, 0, m) {
      B_2[(i*stride)] = ((float32*)B_2[(i*stride)] + (float32*)A_2[((i*stride_1) + (k*stride_2))])
    }
  }
}
"""

也可以查看build之后的LLVM代码：
print(m.get_source())

2. 安装和体验TVM

1. clone代码

• 拉取仓库

  git clone --recursive https://github.com/apache/tvm tvm
  

• 拉取子仓库

  git submodule init
  git submodule update
  

2. docker镜像

• 拉取镜像

  docker pull tvmai/ci-gpu
  或者
  docker pull tvmai/ci-cpu
  

• 启动容器

  cd tvm
  ./docker/bash.sh tvmai/ci-gpu
  

3. 编译TVM

• 编译命令

  mkdir build
  cd build
  cp ../cmake/config.cmake .
  cmake ..
  make -j$(nproc)
  

4. 配置环境变量

```bash
export TVM_HOME=/workspace/tvm
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}
```