# 前言

最近开始学习 TVM。感觉 TVM 英文文档中 TVM Codebase Walkthrough by Example 一节对于理解 TVM 工程非常有用。本篇文章只是翻译，可以直接跳转查看英文全文。

这个时代有这么多愿意开源并将技术介绍给我们的行业大牛，真是我们的荣幸，膜拜！
------ 大家好啊 我是 暮冬 Z 羡慕

# Codebase Structure Overview

在 TVM 存储库的根目录中，我们有以下子目录，它们共同构成了大部分代码库。

• src
  C++ code for operator compilation and deployment runtimes.
  算子编译 、 runtime 部署 的 C++ 代码
• src/relay
  Implementation of Relay, a new functional IR for deep learning framework.
  Relay IR 的实现 算子的映射关系在 src/relay/op
• python
  Python frontend that wraps C++ functions and objects implemented in src.
  python 前端
• src/topi
  Compute definitions and backend schedules for standard neural network operators.
  计算标准神经网络算子的定义和后端调度

TVM 中 Python 和 C++ 的互操作性不是单向的。尽管在 TVM 中 C++ 完成繁重的内部执行工作，Python 完成用户接口， TVM 中也存在 C++ 调用 Python 的情况：For example, the convolution operator is implemented in Python, and its implementation is invoked from C++ code in Relay.（Relay 中的 C++ 调用 Python 实现的卷积算子）

# Vector Add Example

使用 vector add 的例子来查看底层 TVM API.

n = 1024

A = tvm.te.placeholder((n,), name='A')

B = tvm.te.placeholder((n,), name='B')

C = tvm.te.compute(A.shape, lambda i: A[i] + B[i], name="C")

这里面 A、B、C 都是 tvm.tensor.Tensor 其 Python 定义位于 python/tvm/te/tensor.py . 支撑的 C++ 定义位于 include/tvm/te/tensor.h 和 src/te/tensor.cc 所有的 Python 类型定义都能找到对应的相同名字的 C++ 定义。

Python 对 C++ 的包装位于 python/tvm/_ffi/ 。

一个 Tensor 包含一个 Operation 类，定义于 python/tvm/te/tensor.py，对应的 C++ 实现位于 include/tvm/te/operation.h 和 src/tvm/te/operation 。 Tensor 是 Operation 类的输出。

我们将输出张量 C 对应的操作传递给 tvm.te.create_schedule() 函数 （来自于 python/tvm/te/schedule.py 。）

s = tvm.te.create_schedule(C.op)

这个函数映射到 C++ 函数 include/tvm/schedule.h 。

inline Schedule create_schedule(Array<Operation> ops) {

  return Schedule(ops);

}

Schedule 包含 Stage 输出 Operation 的集合。

Stage 对应于一个操作 Operation 。上面的 vector add 操作中有两个 placeholder ops 和一个 compute op. 所以 Schedule s 有三个状态 Stage ，每个 Stage 持有以下信息： 循环嵌套结构、每个循环的类型（ Parallel，Vectorized，Unrolled ）、以及在下一个循环嵌套 Stage 中在哪里执行它自己的计算。

Schedule 和 Stage 本身定义在 tvm/python/te/schedule.py ， include/tvm/te/schedule.h ， 和 src/te/schedule/schedule_ops.cc 。

为简单起见，我们使用 tvm.build(...) 处理上方 create_schedule() 函数创建的默认 Schedule s 和 <em>。我们必须添加必要的线程绑定，来使得其能在 GPU 上运行：

target = "cuda"

bx, tx = s[C].split(C.op.axis[0], factor=64)

s[C].bind(bx, tvm.te.thread_axis("blockIdx.x"))

s[C].bind(tx, tvm.te.thread_axis("threadIdx.x"))

fadd = tvm.build(s, [A, B, C], target)

tvm.build(...) ，定义在 python/tvm/driver/build_module.py ， 需要输入一个 Schedule ; input , output Tensor ; 以及一个 target 。返回一个 tvm.runtime.Module 。

整个 tvm.build(...) 过程可以分成两步：

• i. 降级 高级的、初始的循环嵌套结构被转换为 最终的、低级的 IR

• ii. 代码生成 low level IR 生成目标机器码

降级是通过 tvm.lower() 函数完成的，它定义在 python/tvm/build\_module.py 。第一，指定绑定推理，一个最初的循环嵌套结构就创建好了。

def lower(sch,

          args,

          name="default_function",

          binds=None,

          simple_mode=False):

   ...

   bounds = schedule.InferBound(sch)

   stmt = schedule.ScheduleOps(sch, bounds)

   ...

边界推断是推断所有循环边界和中间缓冲区大小的过程。如果你的目标是 CUDA，且你用了 share memory，它需要的最小 size 在此处确定。绑定推理时在 src/te/schedule/bound.cc，src/te/schedule/graph.cc 和 src/te/schedule/message\_passing.cc 中实现的。

stmt ， ScheduleOps() 的输出，表示一个初识的循环嵌套结构。如果在 schedule 中已经应用了 reorder 和 split 原语，那么初始的循环嵌套结构已经反映了这些变化。 ScheduleOps() 定义在 rc/te/schedule/schedule_ops.cc 。

接下来应用一些 lowering passes to stmt . 这些 passes 在 src/tir/pass 子文件夹下实现。举个例子，如果在你的 schedule 中应用了 vectorize 或者 unroll 原语，他们会被应用到循环 vectorization 和 unrolling passes。

...

stmt = ir_pass.VectorizeLoop(stmt)

...

stmt = ir_pass.UnrollLoop(

    stmt,

    cfg.auto_unroll_max_step,

    cfg.auto_unroll_max_depth,

    cfg.auto_unroll_max_extent,

    cfg.unroll_explicit)

...

在降级 lowering 结束后， build() 函数生成目标机器代码。如果你的设备是 X86, 这个代码可能包含 SSE 或者 AVX 指令；如果是 CUDA 设备，将包含 PTX 指令。 此外，除了目标特定的机器代码之外，TVM 还生成负责内存管理、内核启动等的主机端代码。

build\_module() 函数完成代码生成，定义在 python/tvm/target/codegen.py 。在 C++ 端代码生成定义在 src/target/codegen 。 build\_module() Python 函数会搜索在 src/target/codegen/codegen.cc 中的 build() 函数。

build() 函数 PackedFunc 注册表中为目标设备查找代码生成器，并调用找到的函数。例如， codegen.build\_cuda 函数注册在 src/codegen/build_cuda_on.cc ，就像这样：

TVM_REGISTER_GLOBAL("codegen.build_cuda")

.set_body([](TVMArgs args, TVMRetValue* rv) {

    *rv = BuildCUDA(args[0]);

  });

上方的 BuildCUDA() 函数使用定义在 src/codegen/codegen_cuda.cc 的 CodeGenCUDA 类，从 lowered IR 生成 CUDA kernel source，并使用 NVRTC 编译 kernel。如果你的目标设备使用 LLVM，包括 X86、ARM、NVPTX 和 AMDGPU，代码可由定义在 src/codegen/llvm/codegen_llvm.cc 的 CodeGenLLVM 来生成。 CodeGenLLVM 将 TVM IR 转换成 LLVM IR，运行一些 LLVM 优化 passes，以及生成目标机器码。

在 src/codegen/codegen.cc 中的 Build() 函数会返回一个 runtime::Module 类，它定义在 include/tvm/runtime/module.h 和 src/runtime/module.cc 。一个 Module 类是一个潜在目标 设备的特定 ModuleNode 的容器。

每个后端都实现一个 ModuleNode 的子类，来添加目标特定的 runtime API 调用。 例如，CUDA 后端在 src/runtime/cuda/cuda_module.cc 实现 CUDAModuleNode 类，来管理 CUDA 驱动 API。上方的 BuildCUDA() 函数用 runtime::Module 包装了 CUDAModuleNode ，并包装到 Python 端。LLVM 后端在 src/codegen/llvm/llvm_module.cc 实现了 LLVMModuleNode ，处理了 JIT 执行和编译代码。其他对应各个后端的 ModuleNode 子类可以在 src/runtime 子文件夹找到。
返回的 module ，可以被认作编译函数和设备 API 的组合，可以被 TVM 的 NDArray objects 调用。

dev = tvm.device(target, 0)

a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)

b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)

c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)

fadd(a, b, c)

output = c.numpy()

在幕后，TVM 会自动分配设备内存并管理内存传输。为了实现这个目标，每个后端都需要继承在 include/tvm/runtime/device_api.h 定义的 DeviceAPI 类，使用设备特定的 API 重写里面的内存管理方法。例如，CUDA 后端在 src/runtime/cuda/cuda_device_api.cc 使用 cudaMalloc 、 cudaMemcpy 实现了 CUDADeviceAPI .

第一次使用 fadd(a, b, c) 调用编译后的模块时，会调用 ModuleNode 的 GetFunction() 方法来获取可用于内核调用的 PackedFunc 。例如，在 src/runtime/cuda/cuda_module.cc CUDA 后端实现了 CUDAModuleNode::GetFunction() 函数如下：

PackedFunc CUDAModuleNode::GetFunction(

      const std::string& name,

      const std::shared_ptr<ModuleNode>& sptr_to_self) {

  auto it = fmap_.find(name);

  const FunctionInfo& info = it->second;

  CUDAWrappedFunc f;

  f.Init(this, sptr_to_self, name, info.arg_types.size(), info.launch_param_tags);

  return PackFuncVoidAddr(f, info.arg_types);

}

PackedFunc 的重载函数 operator() 会被调用。从而会调用定义在 src/runtime/cuda/cuda_module.cc 的 CUDAWrappedFunc 的 operator() 函数，最终我们会看到 cuLaunchKernel 驱动会调用：

class CUDAWrappedFunc {

 public:

  void Init(...)

  ...

  void operator()(TVMArgs args,

                  TVMRetValue* rv,

                  void** void_args) const {

    int device_id;

    CUDA_CALL(cudaGetDevice(&device_id));

    if (fcache_[device_id] == nullptr) {

      fcache_[device_id] = m_->GetFunc(device_id, func_name_);

    }

    CUstream strm = static_cast<CUstream>(CUDAThreadEntry::ThreadLocal()->stream);

    ThreadWorkLoad wl = launch_param_config_.Extract(args);

    CUresult result = cuLaunchKernel(

        fcache_[device_id],

        wl.grid_dim(0),

        wl.grid_dim(1),

        wl.grid_dim(2),

        wl.block_dim(0),

        wl.block_dim(1),

        wl.block_dim(2),

        0, strm, void_args, 0);

  }

};

本文概括了 TVM 如何编译和执行函数。 虽然本文没有详细说明 TOPI 或 Relay，但最终所有神经网络算子都会经历与上述相同的编译过程。

# 后记

本博客目前以及可预期的将来都不会支持评论功能。各位大侠如若有指教和问题，可以在我的 github 项目 或随便一个项目下提出 issue，并指明哪一篇博客，我看到一定及时回复！