本文翻译自 TVM 官方文档:TVM Codebase Walkthrough by Example
TVM 代码库示例讲解
了解一个新的代码库可能是一个挑战。 对于像 TVM 这样的代码库尤其如此,其中不同的组件以不明显的方式交互。 在本指南中,我们尝试通过一个简单的示例来说明构成编译管道的关键元素。 对于每个重要步骤,我们都会展示它在代码库中的实现位置。 目的是让新开发人员和感兴趣的用户更快地熟悉代码库。
代码库结构概述
在 TVM 存储库的根目录中,我们有以下子目录,它们共同构成了大部分代码库。
src
- 用于算子编译和部署运行时的 C++ 代码。src/relay
- Relay 的实现,一种用于深度学习框架的新功能 IR。python
- 包装 C++ 函数和在src
中实现的对象的 Python 前端。src/topi
- 计算标准神经网络算子的定义和后端 schedule。
使用标准的深度学习术语,src/relay
是管理计算图的组件,并且图中的节点使用在 src
的其余部分中实现的基础架构进行编译和执行。 python
为用户可以用来执行编译的 C++ API 和驱动程序代码提供 python 绑定。每个节点对应的算子都注册在 src/relay/op
中。运算符的实现在 topi
中,它们是用 C++ 或 Python 编码的。
当用户通过relay.build(...)
调用图编译时,图形中的每个节点都会发生以下一系列操作:
- 通过查询算子注册表,查找算子的实现
- 为算子生成计算表达式和 schedule
- 将算子编译成目标代码
TVM 代码库的有趣方面之一是 C++ 和 Python 之间的互操作性不是单向的。通常,所有执行繁重工作的代码都是用 C++ 实现的,并且为用户界面提供了 Python 绑定。在 TVM 中也是如此,但在 TVM 代码库中,C++ 代码也可以调用 Python 模块中定义的函数。例如,卷积算子是用 Python 实现的,它的实现是从 Relay 中的 C++ 代码调用的。
向量加法示例
我们使用一个直接使用低级 TVM API 的简单示例。示例是向量加法,在 Working with Operators Using Tensor Expression 中有详细介绍。
1 | n = 1024 |
这里,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
的子类。
1 |
|
对象协议是将 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.py
、include/tvm/te/operation.h
和 src/tvm/te/operation
子目录中。Tensor
是其 Operation
对象的输出。每个 Operation
对象依次具有 input_tensors()
方法,该方法返回一个输入 Tensor
列表给它。这样我们就可以跟踪 Operation
之间的依赖关系。
我们将输出 Tensor C
对应的操作传递给 python/tvm/te/schedule.py
中的 tvm.te.create_schedule()
函数。
1 | s = tvm.te.create_schedule(C.op) |
此函数映射到 include/tvm/schedule.h
中的 C++ 函数。
1 | inline Schedule create_schedule(Array<Operation> ops) { |
Schedule
是由 Stage
和输出 Operation
组成的集合。
Stage
对应一个 Operation
。在上面的向量添加示例中,有两个占位符操作和一个计算操作,因此 schedule 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()
函数创建的默认 schedule 上调用 tvm.build(...)
。
1 | target = "cuda" |
tvm.build()
,在 python/tvm/driver/build_module.py
中定义,接受一个 schedule、输入和输出 Tensor 和一个 target,并返回一个 tvm.runtime.Module
对象。 tvm.runtime.Module
对象包含一个已编译的函数,可以使用函数调用语法调用该函数。
tvm.build()
的过程可以分为两个步骤:
-
Lowering,将高级初始循环嵌套结构转换为最终的低级 IR
-
代码生成,其中目标机器代码是从低级 IR 生成的
Lowering 由 tvm.lower()
函数完成,在 python/tvm/build_module.py
中定义。首先,执行绑定推理,并创建初始循环嵌套结构。
1 | def lower(sch, |
边界推断是推断所有循环边界和中间缓冲区大小的过程。如果您以 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
应用许多 lowering pass。这些 pass 在 src/tir/pass
子目录中实现。例如,如果您已将 vectorize
或 unroll
原语应用于您的 schedule,则它们将应用于下面的循环矢量化和展开通道。
1 | ... |
Lowering 完成后, build()
函数从已 lower 的函数生成目标机器代码。如果您针对 x86,此代码可以包含 SSE 或 AVX 指令,或者针对 CUDA 目标包含 PTX 指令。除了目标特定的机器代码,TVM 还生成负责内存管理、内核启动等的主机端代码。
代码生成由 python/tvm/target/codegen.py
中定义的 build_module()
函数完成。在 C++ 端,代码生成在 src/target/codegen
子目录中实现。Python 函数 build_module()
将调用 src/target/codegen/codegen.cc
中的 Build()
函数:
Build()
函数在 PackedFunc 注册表中查找给定目标的代码生成器,并调用找到的函数。例如,在 src/codegen/build_cuda_on.cc
中注册了 codegen.build_cuda
函数,如下所示:
1 | TVM_REGISTER_GLOBAL("codegen.build_cuda") |
上面的BuildCUDA()
使用src/codegen/codegen_cuda.cc
中定义的CodeGenCUDA
类从降低的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()
函数返回一个runtime::Module
对象,在include/tvm/runtime/module.h
和src/runtime/module.cc
中定义. 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 对象上调用。
1 | dev = tvm.device(target, 0) |
在底层,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()
,如下所示:
1 | PackedFunc CUDAModuleNode::GetFunction( |
PackedFunc
的重载 operator()
将被调用,进而调用 src/runtime/cuda/cuda_module.cc
中的 CUDAWrappedFunc
的 operator()
,最后我们看到 cuLaunchKernel
驱动调用:
1 | class CUDAWrappedFunc { |
到此结束对 TVM 如何编译和执行函数的概述。虽然我们没有详细说明 TOPI 或 Relay,但最终所有的神经网络算子都会经历与上面相同的编译过程。我们鼓励您深入了解代码库其余部分的细节。