# 前言
最近开始学习 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，或者知乎 私信，并指明哪一篇博客，我看到一定及时回复，感激不尽！
