TVM代码库演练示例
- 代码库结构概述
- 向量加法示例

TVM代码库演练示例

了解新的代码库可能是一个挑战。尤其是TVM的代码库，其中不同的组件以非显而易见的方式交互。在本指南中，我们尝试通过一个简单的示例来说明构成编译pipeline的关键元素。对于每个重要步骤，我们都会说明在代码库中的哪个位置。目的是让新开发人员和感兴趣的用户更快地深入代码库。

代码库结构概述

在TVM 代码库的根目录中，有以下子目录，这些子目录一起构成了大部分代码库。

src - c++代码，用于算子(operator)编译和部署运行时。
src/relay - Relay的实现，Relay是深度学习框架的一个非常有用的IR。
python- Python前端，封装了在src中实现的c++函数和对象。
topi - 标准神经网络算子(operator)的计算定义和后端调度。

使用标准的深度学习术语，src/relay是管理计算图( graph )的组件，并且图中的节点(node)是使用在src的其余部分(除relay之外)中实现的基础结构来编译和执行的。python为c++ API和驱动程序代码提供python绑定，用户可用来执行编译。每个节点(node)对应的算子(operator)在src/relay/op中注册。算子(operator)在topi中实现，并且使用c++或Python进行编码。
当用户通过relay.build(…)调用图编译时，图中的每个节点(node)都会发生以下操作顺序：

通过查询算子(operator)注册表来查找算子(operator)实现
为算子(operator)生成计算表达式(compute)和调度(schedule)
将算子(operator)编译为目标代码

TVM代码库有趣之一是c++和Python之间的互操作性不是单向的。通常，所有执行繁重工作的代码都是用c++实现的，并且为用户接口提供了Python绑定。但是在TVM代码库中，c++代码也可以调用Python模块中定义的函数。例如，卷积算子(operator)是用Python实现的，但从Relay中的c++代码来调用。

向量加法示例

我们使用一个直接使用低级TVM API的简单示例。该示例是向量加法，本教程将对此进行详细介绍。

n = 1024
A = tvm.placeholder((n,), name='A')
B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")

在这里，A，B，C的类型是tvm.tensor.Tensor(在python/tvm/tensor.py定义)。Python Tensor由c++Tensor支持，c++Tensor在include/tvm/tensor.h和src/lang/tensor.cc中实现。TVM中的所有Python类型都可以视为具有相同名称的基础c++类型的句柄(handle)。如果您查看下面的Python Tensor类型的定义，可以看到它是NodeBase的子类。

@register_node
class Tensor(NodeBase, _expr.ExprOp):
"""Tensor object, to construct, see function.Tensor"""
    def __call__(self, *indices):
...

Node系统是将c++类型公开给前端语言（包括Python）的基础。TVM实现Python包装的方式并不简单。本文档简要介绍了它，如果您有兴趣，详细信息请参阅python/tvm/_ffi/。
Tensor是由python/tvm/api.py中的函数创建的，而这些函数又调用了src/api/api_lang.cc中公开的c++函数。在src/api子目录中公开了可从Python调用的所有c++函数。例如，上面的tvm.compute()函数调用了在src/api/api_lang.cc中公开的API _ComputeOp：

TVM_REGISTER_API("_ComputeOp")
.set_body([](TVMArgs args, TVMRetValue* ret) {
    *ret = ComputeOpNode::make(args[0],
                                args[1],
                                args[2],
                                args[3],
                                args[4]);
});

我们使用TVM_REGISTER_*宏以PackedFunc的形式将c++函数公开给前端语言。 PackedFunc是TVM在c++和Python之间实现互操作性的另一种机制。特别是，这使得从c++代码库调用Python函数非常容易。
一个Tensor对象具有与其相关联的Operation对象(在python/tvm/tensor.py，include/tvm/operation.h和src/tvm/op子目录定义)。Tensor是其Operation对象的输出。每个Operation对象都有相应的input_tensors()方法，该方法返回一个输入Tensor列表。这样我们就可以追踪Operation之间的依赖关系。
我们将与输出张量C相对应的运算传递给函数tvm.create_schedule()(在python/tvm/schedule.py中定义)。
s`` ``=`` ``tvm``.``create_schedule``(``C``.``op``)
此函数映射到include/tvm/schedule.h中的c++函数。

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

Schedule由Stage和输出Operation的集合组成。
Stage对应一个Operation。在上面的矢量添加示例中，有两个占位符操作和一个计算操作，因此调度s包含三个阶段。各Stage保持关于循环嵌套结构的信息、各循环(Parallel，Vectorized，Unrolled)的类型、以及在下一步的Stage中的循环嵌套的哪个地方执行它的计算（如果有的话）。
Schedule和Stage在tvm/python/schedule.py，include/tvm/schedule.h和src/schedule/schedule_ops.cc中定义。
为简单起见，我们在上述函数create_schedule()创建的默认调度中调用tvm.build(…)。

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

tvm.build()（在python/tvm/build_module.py中定义）获取调度、输入和输出Tensor以及目标，然后返回在python/tvm/module.py中定义的tvm.Module对象。一个Module对象包含一个已编译函数，该函数可以通过函数调用语法来调用。
tvm.build()的过程可以分为两个步骤：

降级(Lowering)，将高级别的初始循环嵌套结构转换为最终的低级别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)
    ...

边界推理(Bound inference)是这样一个过程，它推断所有循环边界以及中间缓冲区的大小。如果您以CUDA后端为目标并使用共享内存，则会在此处自动确定其所需的最小大小。边界推理在src/schedule/bound.cc，src/schedule/graph.cc和src/schedule/message_passing.cc中实现。有关边界推理如何工作的更多信息，请参见InferBound Pass。
stmt是ScheduleOps()的输出，代表初始的循环嵌套结构。如果您已将原语reorder或split应用到你的调度，则初始循环嵌套已经反映了这些更改。ScheduleOps()在src/schedule/schedule_ops.cc中定义。
接下来，我们将多个降级pass应用于stmt。这些pass在src/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)
...

降级完成后，build()函数从降级的函数生成目标机器码。如果您以x86为目标，则此代码可以包含SSE或AVX指令，或以CUDA为目标的PTX指令。除了特殊目标的机器码外，TVM还生成主机端代码，该代码负责内存管理、内核启动等。
代码生成由build_module()函数（python/tvm/codegen.py中定义）完成。在c++方面，代码生成在子目录src/codegen中实现。Python函数build_module()将调用src/codegen/codegen.cc 中的Build()函数，如下所示：

runtime::Module Build(const Array<LoweredFunc>& funcs,
const std::string& target) {
    std::string build_f_name = "codegen.build_" + target;
    const PackedFunc* bf = runtime::Registry::Get(build_f_name);
    runtime::Module m = (*bf)(funcs, target);
    return m;
}

该Build()函数在PackedFunc注册表中查找给定目标的代码生成器，并调用找到的函数。例如，codegen.build_cuda函数在src/codegen/build_cuda_on.cc中注册，如下所示：

TVM_REGISTER_API("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优化pass，并生成目标机器码。
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()函数用runtime::Module包装CUDAModuleNode，并返回到Python端。LLVM后端在src/codegen/llvm/llvm_module.cc中实现LLVMModuleNode，它处理已编译代码的JIT执行。ModuleNode的其他子类可以在src/runtime与各后端相对应的子目录下找到。
返回的模块（可以认为是已编译函数和设备API的组合）可以在TVM的NDArray对象上调用。

ctx = tvm.context(target, 0)
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
fadd(a, b, c)
output = c.asnumpy()

在后端，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.thread_axis_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 = thread_axis_cfg_.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。最后，所有神经网络算子(operator)都经过与上述相同的编译过程。欢迎您深入研究代码库的其它细节。

人工智能

[翻译]TVM代码库演练示例— tvm 0.6.0文档

TVM代码库演练示例

代码库结构概述

向量加法示例