TVM Developer Tutorial

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 的简单示例。 示例是向量加法,在使用张量表达式处理运算符中有详细介绍

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 张量由 C++ 张量支持,在 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):
       ...

对象协议是将 C++ 类型暴露给前端语言(包括 Python)的基础。 TVM 实现 Python 包装的方式并不简单。 TVM Runtime 系统中简要介绍了它,如果您有兴趣,详细信息在 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.hsrc/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。 在上面的向量添加示例中,有两个占位符操作和一个计算操作,因此调度 s 包含三个阶段。 每个 Stage 保存有关循环嵌套结构、每个循环的类型(ParallelVectorizedUnrolled)以及在下一个 Stage 的循环嵌套中执行其计算的位置(如果有)的信息。
ScheduleStagetvm/python/te/schedule.pyinclude/tvm/te/schedule.hsrc/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 中定义,接受一个时间表、输入和输出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)
   ...

边界推断是推断所有循环边界和中间缓冲区大小的过程。 如果您以 CUDA 后端为目标并使用共享内存,则在此处自动确定其所需的最小大小。 绑定推理在 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 应用许多lowering passes。 这些passes在 src/tir/pass 子目录中实现。 例如,如果您已将矢量化或展开原语应用于您的日程安排,则它们将应用于下面的循环矢量化和展开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 注册表中查找给定目标的代码生成器,并调用找到的函数。 例如,在 src/codegen/build_cuda_on.cc 中注册了 codegen.build_cuda 函数,如下所示:

TVM_REGISTER_GLOBAL("codegen.build_cuda")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    *rv = BuildCUDA(args[0]);
  });

上面的 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.hsrc/runtime/module.cc 中。 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 自动分配设备内存并管理内存传输。 为此,每个后端都需要继承在 include/tvm/runtime/device_api.h 中定义的 DeviceAPI 类,并覆盖内存管理方法以使用特定于设备的 API。 例如,CUDA 后端在 src/runtime/cuda/cuda_device_api.cc 中实现 CUDADeviceAPI 以使用 cudaMalloccudaMemcpy 等。

第一次使用 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() 将被调用,进而调用 src/runtime/cuda/cuda_module.ccCUDAWrappedFuncoperator(),最后我们看到 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,但最终所有的神经网络算子都会经历与上面相同的编译过程。 我们鼓励您深入了解代码库其余部分的细节。

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值