源碼研習 — 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}
```