通过示例逐步讲解 TVM 代码库

通过示例逐步讲解 TVM 代码库#

熟悉新的代码库可能是一项挑战。对于像 TVM 这样的代码库来说尤其如此，因为它的不同组件以非显而易见的方式相互作用。在本指南中，尝试通过简单的示例来说明构成编译管线的关键元素。对于每个重要步骤，展示了它在代码库中的实现位置。目的是让新开发者和感兴趣的用户能够更快地深入代码库。

代码库结构概览#

在 TVM 仓库的根目录下，有以下几个子目录，它们共同构成了代码库的主要部分。

  • src - 用于算子编译和部署运行时的 C++ 代码。

  • src/relay - Relay 的实现，这是为深度学习框架设计的新型函数式中间表示（IR）。

  • python - Python 前端，封装了在 src 中实现的 C++ 函数和对象。

  • 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 的简单示例。该示例是向量加法，在 tutorial-tensor-expr-get-started 中有详细介绍。

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 支持，后者实现在 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运行时系统 中有简要介绍，如果你感兴趣，细节可以在 python/tvm/_ffi/ 中找到。

使用 TVM_REGISTER_* 宏将 C++ 函数以 PackedFunc 的形式暴露给前端语言。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() 方法，该方法返回其输入 Tensor 的列表。通过这种方式，可以跟踪 Operation 之间的依赖关系。

将与输出张量 C 对应的运算传递给 python/tvm/te/schedule.py 中的 tvm.te.create_schedule() 函数。

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

此函数映射到 include/tvm/schedule.h 中的 C++ 函数。

inline Schedule create_schedule(Array<Operation> ops) {
  return Schedule(ops);
}

Schedule 由 Stage 集合和输出 Operation 组成。

Stage 对应一个 Operation。在上面的向量加法示例中，有两个占位符操作和一个计算操作，因此调度 s 包含三个阶段。每个 Stage 都保存了关于循环嵌套结构的信息、每个循环的类型（Parallel、Vectorized、Unrolled），以及如果有下一个 Stage，则在其循环嵌套中执行计算的位置。

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

为了简化，在上述 create_schedule() 函数创建的默认调度上调用 tvm.build(...)，并且必须添加必要的线程绑定以使其可在 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 中，它接收调度、输入和输出 Tensor 以及目标，并返回 tvm.runtime.Module 对象。tvm.runtime.Module 对象包含编译后的函数，可以通过函数调用语法来调用它。

tvm.build() 的过程可以分为两个步骤：

  • Lowering，将高级的初始循环嵌套结构转换为最终的低级 IR。

  • 代码生成，从低级 IR 生成目标机器代码。

Lowering 由定义在 python/tvm/build_module.py 中的 tvm.lower() 函数完成。首先，执行边界推断，并创建初始的循环嵌套结构。

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 中。

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

接下来，对 stmt 应用一系列 lowering 过程。这些过程实现在 src/tir/pass 子目录中。例如，如果你对调度应用了 vectorize 或 unroll 原语，它们将在下面的循环向量化和展开过程中应用。

...
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() 函数从 lowered 函数生成目标机器代码。如果你的目标是 x86，此代码可能包含 SSE 或 AVX 指令；如果是 CUDA 目标，则可能包含 PTX 指令。除了目标特定的机器代码外，TVM 还会生成负责内存管理、内核启动等的主机端代码。

代码生成由定义在 python/tvm/target/codegen.py 中的 build_module() 函数完成。在 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 内核源代码，并使用 NVRTC 编译内核。如果你的目标是使用 LLVM 的后端（包括 x86、ARM、NVPTX 和 AMDGPU），代码生成主要由定义在 src/codegen/llvm/codegen_llvm.cc 中的 CodeGenLLVM 类完成。CodeGenLLVM 将 TVM IR 转换为 LLVM IR，运行一系列 LLVM 优化过程，并生成目标机器代码。

src/codegen/codegen.cc 中的 Build() 函数返回定义在 include/tvm/runtime/module.h 和 src/runtime/module.cc 中的 runtime::Module 对象。Module 对象是底层目标特定的 ModuleNode 对象的容器。每个后端都实现了 ModuleNode 的子类以添加目标特定的运行时 API 调用。例如，CUDA 后端在 src/runtime/cuda/cuda_module.cc 中实现了 CUDAModuleNode 类，用于管理 CUDA 驱动 API。上述的 BuildCUDA() 函数将 CUDAModuleNode 包装在 runtime::Module 中并返回给 Python 端。LLVM 后端在 src/codegen/llvm/llvm_module.cc 中实现了 LLVMModuleNode，用于处理编译代码的 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 会自动分配设备内存并管理内存传输。为此，每个后端都需要子类化定义在 include/tvm/runtime/device_api.h 中的 DeviceAPI 类，并重写内存管理方法以使用设备特定的 API。例如，CUDA 后端在 src/runtime/cuda/cuda_device_api.cc 中实现了 CUDADeviceAPI，以使用 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);
}

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，但最终所有神经网络算子都会经历与上述相同的编译过程。鼓励你深入研究代码库其余部分的细节。