【TVM帮助文档学习】通过示例分析TVM代码流程

本文翻译自TVM Codebase Walkthrough by Example — tvm 0.9.dev0 documentation

了解TVM代码颇具挑战性，它的各组件之间的交互非常隐晦。 在本指南中，我们将通过一个简单的示例来说明模型编译过程中的关键组成。 对于每一个重要的步骤，我们都展示了它在代码库中的实现位置，以便新开发人员和感兴趣的用户更快地深入代码库。

代码结构概述

在TVM代码仓库根目录下有如下子目录，它们组成了TVM代码库：

• src - 算子编译和部署运行时C++代码
• src/relay - Relay（一种用于深度学习框架的函数式IR）的实现
• python - src中C++函数和对象的Python封装前端.
• src/topi - 标准神经网络算子的计算定义和后端调度.

按照标准的深度学习术语， src/relay是管理计算图的组件。计算图中的节点由src目录下其他子目录提供的基础设施来编译和执行。python为c++ API和驱动程序代码提供了python绑定，用户可以使用这些绑定来执行编译。各个节点对应的操作符注册在src/relay/op中。算子的实现是在topi目录下，它们是用c++或Python实现的。

当用户通过relay.build(…)编译图时，图中的各节点会执行以下一系列操作:

• 通过查询运算符注册表来查找算子实现
• 为算子生成计算表达式和调度
• 将操作符编译为目标代码

TVM代码库中一个有趣的地方是，c++和Python之间的调用不是单向的。通常，所有执行繁重任务的代码都用c++实现，而Python绑定提供面向用户的接口。在TVM中也是如此，但在TVM代码库中，c++代码也可以调用Python模块中定义的函数。例如，卷积算子在Python中实现，在Relay的c++代码中调用。

向量加示例

 下面的示例直接使用低级TVM API实现。这个例子是向量加法，在Working with Operators Using Tensor Expression中有详细介绍。

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/tvm/te/tensor.py中。Python Tensor是对C++ Tensor的封装，C++ Tensor的实现在include/tvm/te/tensor.h和src/te/tensor.cc中。TVM中的所有Python类型都可以看作是底层C++同名类型的句柄。如果你看看下面Python Tensor类型的定义，你会发现它是Object的一个子类。

@register_object
class Tensor(Object, _expr.ExprOp):
    """Tensor object, to construct, see function.Tensor"""

    def __call__(self, *indices):
       ...

对象协议是将C++类型公开给前端语言(包括Python)的基础。TVM的Python封装实现方式并不简单。TVM Runtime System中做了简要介绍，如果你感兴趣，可以在python/tvm/_ffi/中找到详细信息。

 我们使用TVM_REGISTER_*宏，以PackedFunc的形式向前端语言公开C++函数。PackedFunc是TVM实现C++和Python互调的另一种机制。特别是，这使得在C++代码中调用Python函数非常容易。你也可以使用FFI Navigator，它能帮你在python和C++的FFI调用之间导航。

Tensor对象与Operation对象关联，定义在python/tvm/te/Tensor.py、include/tvm/te/operation.h和src/tvm/te/operation子目录中。Tensor是Operation对象的输出。每个Operation对象又有一个input_tensors()方法，该方法向Operation返回一个输入张量的列表。这样我们就可以跟踪Operation之间的依赖关系。

下面的代码中，我们将输出张量C的操作传递给tvm.te.create_schedule()。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。在上面的向量加示例中，有两个占位符操作和一个计算操作，因此schedule包含三个Stage。每个Stage保存相关的循环嵌套结构、每个循环的类型(并行、向量化、展开)、以及在（下一个Stage的循环嵌套中的）何处执行当前Stage的计算等信息。

Schedule和Stage定义在tvm/python/te/schedule.py、include/tvm/te/schedule.h和src/te/schedule/schedule_ops.cc中。

为了简单起见，我们对上面create_schedule()函数创建的默认调度调用tvm.build(…)。

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

tvm.build()（定义在python/tvm/driver/build_module.py中）接收一个schedule、输入和输出Tensor，以及一个target，并返回一个tvm.runtime.Module对象。tvm.runtime.Module对象包含一个可以像函数一样调用的编译过的函数。

tvm.build()的处理分两个阶段：

• lower。一个高层的初始的循环嵌套结构转变为最终的底层的IR
• 代码生成。由底层IR生成目标机代码

Lowering由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后端并且使用共享内存，缓存大小下限将在这里自动确定。绑定推断在src/te/schedule/ Bound.cc, src/te/schedule/graph.cc和src/te/schedule/message_passing.cc中实现。有关绑定推断如何工作的更多信息，请参阅InferBound Pass

stmt是ScheduleOps()的输出，表示初始循环嵌套结构。如果你对调度应用了reorder或split原语，那么初始循环嵌套已经反映了这些变更。ScheduleOps()定义在src/te/schedule/schedule_ops.cc中。

接下来，我们对stmt应用一些lower pass。这些pass在src/tir/pass子目录中实现。例如，如果您已经对调度应用了vectorize或unroll原语，那么它们将在下面的循环向量化和循环展开pass中应用。 

..
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)
...

lower完成后，build()函数为lower后的函数生成目标机器代码。目标代码可以包含SSE或AVX指令（如果你的目标是x86），或PTX指令（如果你的目标是CUDA）。TVM还生成特定的host侧代码，包括内存管理、内核启动等。

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()使用CodeGenCUDA类为lower IR生成CUDA内核代码，并使用NVRTC编译代码。CodeGenCUDA定义在src/codegen/codegen_cuda.cc中。如果你的后端使用LLVM（支持x86, ARM, NVPTX和AMDGPU），代码生成主要由定义在src/codegen/ LLVM /codegen_llvm.cc中的CodeGenLLVM类来完成。CodeGenLLVM将TVM IR转换为LLVM IR，执行多个LLVM优化pass，然后生成目标机器代码。

Build()函数（定义见src/codegen/codegen.cc）返回一个runtime::Module对象（定义见include/tvm/runtime/Module.h和src/runtime/Module .cc）。Module对象是指定底层目标ModuleNode对象的容器。每个后端实现一个ModuleNode的子类，以添加指定目标的运行时API调用。例如，CUDA后端实现CUDAModuleNode类（定义见src/runtime/cuda/cuda_module.cc），它管理CUDA驱动程序API。上面的BuildCUDA()函数使用runtime::Module封装CUDAModuleNode，并将其返回给Python端。LLVM后端实现LLVMModuleNode（定义见src/codegen/llvm/llvm_module.cc），它处理已编译代码的JIT执行。ModuleNode的其他子类可以在src/runtime下找到，分别对应对应于各种后端

返回的模块可以看作是一个编译函数和一个设备API的组合，可以在TVM的NDArray对象上调用。

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自动分配设备内存并管理内存转换。为此每个后端定义自己的DeviceAPI(在include/tvm/runtime/device_api.h中定义)子类，并提供内存管理方法以使用对应设备的API。例如，CUDA后端实现CUDADeviceAPI（见src/runtime/ CUDA /cuda_device_api.cc）以使用cudaMalloc, cudaMemcpy等接口。

当你第一次用有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);
}

 packkedfunc的重载操作符()被调用，进而调用(src/runtime/cuda/cuda_module.cc中)CUDAWrappedFunc的重载操作符()。最后我们看到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，但最后所有的神经网络算子都经过了与上述相同的编译过程。我们鼓励您深入研究代码库其余部分的细节。