计算库层

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。

其中：

IR（Intermediate Representation）：一种中间表示形式，用于编译过程中的高级代码表示。
算子（Operator）：在深度学习中，算子通常指代执行特定计算的函数，比如卷积、矩阵乘等。
调度（Schedule）：定义了算子如何在硬件上执行的策略，包括循环的嵌套结构、并行化、向量化等。

向量加法示例：

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")

在 python/tvm/te/tensor.py 中定义的 A、B 和 C 的类型都是 tvm.tensor.Tensor。这些 Python Tensor 是由 C++ Tensor 支持的，其实现位于 include/tvm/te/tensor.h 和 src/te/tensor.cc 文件中。在 TVM 中，所有的 Python 类型都可以视为与其底层 C++ 类型具有相同名称的句柄。

从以下 Python Tensor 类型的定义中可以看出，tvm.tensor.Tensor 是 Object 的一个子类。

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

在 TVM 中，每个 Tensor 对象都有一个与之关联的 Operation 对象。Tensor 是在计算过程中存储数据的多维数组，而 Operation 表示对一个或多个 Tensor 进行操作的计算。这两个概念在代码中有明确的实现，相关定义分别在 python/tvm/te/tensor.py、include/tvm/te/operation.h 和 src/tvm/te/operation 目录下。
每个 Tensor 对象都可以看作是其相应的 Operation 的输出，这意味着通过执行某个 Operation 可以生成一个 Tensor。
Operation 对象提供了一个 input_tensors() 方法，这个方法返回一个输入 Tensor 的列表。这使得开发者能够跟踪不同 Operation 之间的依赖关系，了解一个 Operation 需要哪些输入 Tensor，以及这些输入 Tensor 是由哪些其他 Operation 产生的。
在计算图中，当我们想要调度某个计算时，需要将输出张量（例如上面提到的 C 张量）对应的 Operation 对象传递给 python/tvm/te/schedule.py 中的 tvm.te.create_schedule() 函数create_schedule() 函数负责生成计算的调度策略，以优化计算的执行。这是构建高效计算图的重要步骤，因为它允许对计算的执行顺序和方式进行控制，从而提高性能。

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

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

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

在 TVM 中，调度由多个 Stage 和输出的 Operation 组成。每个 Stage 代表一个 Operation 的计算过程。
以“向量加法”（Vector Add）为例，假设有两个占位符 Operation 和一个计算 Operation，那么这个调度（s）将包含三个阶段（Stage）。
每个 Stage 存储有关循环嵌套的信息，包括：循环嵌套结构：描述了如何将计算划分为多个循环的结构。循环类型：标识每个循环的执行方式，比如：Parallel（并行）：表示该循环可以在多个线程中并行执行。Vectorized（向量化）：表示该循环将数据分块处理，以提高效率。Unrolled（展开）：表示将循环展开为多个相同的操作，以减少循环开销。位置：指明在下一个 Stage 的循环嵌套中执行该计算的位置（如果有嵌套的话）。create_schedule() 函数的作用：create_schedule() 函数用于创建默认的调度。这个调度提供了基础的计算顺序和结构。默认的调度通常会调用 tvm.build(...) 函数来生成可执行的代码。
为了使调度可以在 GPU 上运行，需要为调度中的 Stage 绑定必要的线程。这一步骤是非常重要的，因为 GPU 的并行计算能力依赖于对线程的有效管理和分配。
通过线程绑定，开发者可以控制计算的并行性，从而充分利用 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)

target = "cuda" 指定了目标平台是CUDA，意味着生成的代码将在GPU上运行。
split和bind是调度操作，用于优化并行执行。split将计算操作的循环分割成更小的部分，bind将这些分割的部分绑定到GPU的线程和块上。
tvm.build函数接受调度、输入和输出Tensor以及目标平台，然后返回一个可以在该平台上运行的模块。

tvm.build() 函数：

tvm.build() 函数定义在 python/tvm/driver/build_module.py 中。它的主要作用是接收一个调度（schedule）、输入和输出的 Tensor，以及一个目标设备（target），然后返回一个 tvm.runtime.Module 对象。返回的 tvm.runtime.Module 对象包含一个可以通过函数调用的已编译函数，这意味着用户可以直接调用这个编译后的函数进行计算，而无需关心底层实现细节。
tvm.build() 的过程可以分为两个主要步骤：降级：降级过程将高级、初始的循环嵌套结构转化为最终的底层中间表示（IR）。这一过程是由 tvm.lower() 函数完成的，tvm.lower() 也定义在 python/tvm/build_module.py 中。降级的第一步是进行边界推断，确定每个循环的迭代范围，以便在生成 IR 时确保计算的正确性。随后，tvm.lower() 会创建一个初始的循环嵌套结构，以便更好地表达计算的逻辑和顺序。代码生成：在降级完成后，接下来的步骤是根据底层的 IR 生成目标机器代码。这一过程涉及将 IR 转换为特定硬件可以理解和执行的机器代码。
降级的过程有助于将更高级的计算抽象（例如高层的循环结构和调度策略）转化为更为底层的表示，使得后续的代码生成过程能够更加有效地针对特定硬件进行优化。通过将计算表示降级到 IR，TVM 能够更灵活地进行优化并适配多种硬件目标

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/te/schedule/bound.cc
src/te/schedule/graph.cc
src/te/schedule/message_passing.cc
这些实现文件负责具体的边界推断算法和逻辑，包括如何根据调度信息推断出循环的边界和缓冲区的大小。

ScheduleOps() 的作用：

stmt 是 ScheduleOps() 函数的输出，表示一个初始的循环嵌套结构。这个结构是调度的基础，反映了计算中循环的组织方式。
如果调度过程中已经应用了 reorder 或 split 等原语，则 stmt 将反映这些变化，确保生成的初始循环结构与应用的调度操作一致。
ScheduleOps() 函数的定义位于 src/te/schedule/schedule_ops.cc 中。

接下来，对 stmt 在 src/tir/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() 函数负责从降级后的函数生成特定目标的机器代码。这一步是将中间表示（IR）转化为实际可执行的代码。
如果目标是 x86 架构，生成的代码将包含 SSE（Streaming SIMD Extensions）或 AVX（Advanced Vector Extensions）指令，以优化计算性能。
如果目标是 CUDA，生成的代码将包含 PTX（Parallel Thread Execution）指令，这是 NVIDIA 的一种中间表示，用于描述并行计算的指令。
除了生成目标专用的机器代码，TVM 还会生成一段宿主机代码。这部分代码负责执行一些重要的任务，如内存管理和内核启动等。宿主机代码确保了生成的内核能够在目标设备上正确运行并管理资源。
代码生成的具体实现是在 build_module() 函数中完成的，该函数定义在 python/tvm/target/codegen.py 中。这个 Python 函数负责协调代码生成的各个环节。
在 C++ 端，代码生成的实现细节位于 src/target/codegen 子目录中。这里包含了许多与代码生成相关的实现和优化。
build_module() 函数最终会调用 C++ 端的 Build() 函数，后者位于 src/target/codegen/codegen.cc 中。Build() 函数负责将具体的代码生成逻辑实现，完成从中间表示到目标机器代码的转换。

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 内核源代码。这意味着它将高层的计算表示转化为适合在 NVIDIA GPU 上执行的 CUDA 代码。
生成的 CUDA 内核源代码随后通过 NVRTC（NVIDIA Runtime Compilation）进行编译。NVRTC 是 NVIDIA 提供的一个库，允许在运行时编译 CUDA 程序，方便动态加载和执行。
如果目标是使用 LLVM 后端（如 x86、ARM、NVPTX 和 AMDGPU），代码生成主要由定义在 src/codegen/llvm/codegen_llvm.cc 中的 CodeGenLLVM 类完成。
CodeGenLLVM 的作用是将 TVM 的 IR 转换为 LLVM 的 IR。这一步是重要的，因为 LLVM 提供了强大的优化和代码生成能力。
在生成 LLVM IR 后，CodeGenLLVM 会执行一些 LLVM 优化。这些优化可以提高生成代码的性能，利用 LLVM 的优化工具链来提升最终机器代码的执行效率。
最后，CodeGenLLVM 会生成适用于特定目标架构的机器代码，使得该代码可以在不同的硬件上高效运行。

AI技术栈解析及应用- 作者：张真瑜 | 山东大学智能创新研究院

计算库层