TVM 代码库的实例演练#

了解新的代码库可能是个挑战。对于像 TVM 这样的代码库尤其如此,其中不同的组件以非明显的方式进行交互。在本指南中,试图通过简单的例子来说明构成编译管道的关键因素。对于每一个重要的步骤,我们都显示了它在代码库中的实现位置。其目的是让新的开发者和感兴趣的用户更快进入代码库。

代码库结构概述#

在 TVM 资源库的根部,我们有以下子目录,它们共同构成了代码库的大部分。

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

  • src/relay - Relay 的实现,新的深度学习框架的函数式 IR。

  • python - Python 前端,包装 src 中实现的 C++ 函数和对象。

  • src/topi - 标准神经网络算子的计算定义和后端调度。

使用标准的深度学习术语,src/relay 是管理计算图的组件,图中的节点使用 src 其他部分实现的基础设施进行编译和执行。python 为 C++ API 和 driver 代码提供 python 绑定,用户可以用它来执行编译。与每个节点对应的算子在 src/relay/op 中注册。算子的实现在 topi 中,它们是用 C++ 或 Python 编码的。

当用户通过 relay.build(...) 调用图的编译时,对图中的每个节点都会发生以下一系列动作:

  • 通过查询算子注册表查找算子实现

  • 为算子生成计算表达式和调度

  • 将算子编译成目标代码

TVM 代码库的一个有趣方面是,C++ 和 Python 之间的互操作性不是单向的。通常情况下,所有执行繁重工作的代码都是用 C++ 实现的,而 Python 绑定是为用户接口提供的。这在 TVM 中也是如此,但在 TVM 代码库中,C++ 代码也可以调用 Python 模块中定义的函数。例如,卷积算子是用 Python 实现的,它的实现是由 Relay 中的 C++ 代码调用的。

向量加法示例#

使用简单的例子介绍如何直接使用低级 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")

这里,ABC 的类型是 tvm.tensor.Tensor,定义在 python/tvm/te/tensor.py。Python 的 Tensor 由 C++ 的 Tensor 支持,在 include/tvm/te/tensor.hsrc/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):
       ...

对象协议(object protocol)是将 C++ 类型暴露给前端语言(包括 Python)的基础。TVM 实现 Python 封装的方式并不直接。TVM 运行时系统 中简要介绍了这一点,如果你有兴趣,细节在 python/tvm/_ffi/ 中。

使用 TVM_REGISTER_* 宏,以 PackedFunc 的形式,将 C++ 函数暴露给前端语言。PackedFunc 是 TVM 实现 C++ 和 Python 之间互操作的另一种机制。特别是,这使得从 C++ 代码库中调用 Python 函数非常容易。你也可以查看 FFI Navigator,它允许你在 Python 和 C++ FFI 调用之间进行导航。

Tensor 对象有与之相关的 Operation 对象,定义在 python/tvm/te/tensor.pyinclude/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);
}

ScheduleStage 和输出 Operation 的集合组成。

Stage 对应一个 Operation。在上面的 vector add 例子中,有两个占位符算子和一个计算算子,所以调度 s 包含三个阶段。每个 Stage 都包含关于循环嵌套结构的信息,每个循环的类型(ParallelVectorizedUnrolled),以及在下一个 Stage 的循环嵌套中执行计算的位置。

ScheduleStage 被定义在 tvm/python/te/schedule.pyinclude/tvm/te/schedule.hsrc/te/schedule/schedule_ops.cc

To keep it simple, we call tvm.build(...) on the default schedule created by create_schedule() function above, and we must add necessary thread bindings to make it runnable on 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() 的过程可以分为两个步骤:

  • 降级,即高级的初始循环嵌套结构变换为最终的低级 IR

  • 代码生成,目标的机器码由低级 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)
   ...

边界推断是推断所有循环边界和中间缓冲区 size 的过程。如果您的目标是 CUDA 后端,并且您使用共享内存,那么它所需的最小 size 将在这里自动确定。边界推断在 src/te/schedule/bound.ccsrc/te/schedule/graph.ccsrc/te/schedule/message_passing.cc 中实现。有关边界推断如何工作的更多信息,请参阅 InferBound Pass

stmtScheduleOps() 的输出)表示初始循环嵌套结构。如果已经将 reordersplit 原语应用到调度中,那么初始的循环嵌套已经反映了这些变化。ScheduleOps() 定义在 src/te/schedule/schedule_ops.cc 中。

接下来,对 stmt 应用一些降级 pass。这些 pass 在 src/tir/pass 子目录中实现。例如,如果已经将 vectorizeunroll 原语应用到调度中,它们将在下面的循环 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)
...

降级完成后,build() 函数从降级的函数生成目标机器码。如果目标是 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() 从降级的 IR 使用定义在 src/codegen/codegen_cuda.cc 中的 CodeGenCUDA 类生成 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 中的 runtime::Module 对象。Module 对象是底层目标专用的 ModuleNode 对象的容器。每个后端实现 ModuleNode 子类,以添加目标专用的运行时 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 的子目录下找到。

返回的模块,可以被认为是编译函数和设备 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 自动分配设备内存并管理内存传输(memory transfer)。

你首次用 fadd(a, b, c) 调用编译过的模块时,ModuleNodeGetFunction() 方法会被调用以获得 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() 将被调用,这反过来调用 CUDAWrappedFunc 中的 operator(),在 src/runtime/cuda/cuda_module.cc,最后看到 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,但最终,所有的神经网络算子都经历了上述相同的编译过程。鼓励您深入研究其余代码库的细节。