TVM的“hello world“基础流程 II

上一篇《TVM的“hello world“基础流程 I》中基于一个最基本的case介绍了TVM中计算的定义与schedule的构建。这篇沿用上一篇中的case，继续介绍接下去的一个重点部分，就是编译。

有了前面构建的schedule之后，接着就需要编译并生成目标代码了。这个工作主要由和两个函数来完成。它们的区别在于应用目标的范围，前者用于单个算子，后者用于整个网络。由于网络可看作由算子组成，后者会调用前者。本例中是针对单个算子的，因此这里使用的是前者：

其中最主要的函数定义在文件中。该函数基于给定参数构建出可调用的目标函数。按照官方介绍里的说法，它主要做两个工作 ：

• Lowering：将high-level的循环嵌套结构转换成最终的low-level的IR。
• Codegen：从low-level的IR生成目标机器代码。

该函数的第一个参数是前面构建出来的schedule，第二个参数是函数的参数列表，第三个参数是target。它提供用于lowering和codegen所需的目标平台信息。代码中对应的对象定义在文件中。其构造函数有两个参数，其中第一个参数target指示目标平台的配置。其中的配置项比如：

• kind: 平台类型，它基本决定了生成的代码是在什么处理器上运行。注册的target kind详细见，有llvm, c, cuda, nvptx, romc, opencl, metal, vulkan, hexagon等。
• keys: 如kind是opencl的话，key可以是mali, opencl, gpu。
• device：对应实际运行的设备，它会添加到keys后面。
• libs：外部库，如cblas, cudnn, cublas, mkl这些。
• …

另外参数host与target类似，但它用于指示host平台。比如taret平台为cuda的话，毕竟GPU还是不能完全脱离CPU运行，因此还需要host的代码做胶水，如内存分配，kernel启动这些。默认为llvm。

Lowering过程可以单独用函数完成，如：

也可以通过函数完成（因为它一进去就会先调用函数）。函数的主要流程相关代码：

它主要根据参数给的schedule与参数生成对应的对象（定义在中）。是软件栈中所有IR变换的基础单元。它维护函数与类型定义。这里的各种pass就是在上进行并吐出。

函数中有四个阶段，第一个阶段中通过函数根据给定的schedule生成对象，然后在这个对象上应用4轮的pass。这些pass主要分为几个阶段，分别是：

• Phase 0：使用者自定义的pass。
• Phase 1：使用者自定义的pass。以及：
  • InjectPrefetch
  • StorageFlatten
  • BF16Legalize
  • NarrowDataType
  • Simplify
• Phase 2：使用者自定义的pass。以及：
  • LoopPartition
  • VectorizeLoop
  • InjectVirtualThread
  • InjectDoubleBuffer
  • StorageRewrite
  • UnrollLoop
• Phase 3：使用者自定义的pass。以及：
  • Simplify
  • RemoveNoOp
  • RewriteUnsafeSelect
  • HoistIfThenElse
  • InstrumentBoundCheckers

这此pass其实是编译构建过程中的精华之一。但限于篇幅（其实是我自己也没了解全。。。），以后再进一步讨论。

函数的最后返回经过上面多轮pass优化后的对象。其中函数是相对比较复杂的一部分，它主要负责生成最初的对象，其中几个关键步骤如下：

1. 函数规范化给定的schedule。主要实现在文件中。它调用以下三个函数。本例比较简单，因此它们实际都没有起什么作用。。。
   1. 函数处理算子内联。用到调度原语 的话会用到。
   2. 函数将循环迭代的最小界置为0。感觉有点canonicalization的意思。
   3. 函数处理在使用调度原语时且目标迭代又被split或fuse情况下的合法化。
2. 函数顾名思义就是边界推导（Bound inference），主要用于推导循环边界。更具体地，就是确定每个的范围，它返回到的映射，即每个循环变量的范围。这个信息在后面的函数中用于确定for循环的范围，和在函数中设置缓冲的大小。具体可参见官方文档 InferBound Pass。
3. 函数基于前面经过一些处理后的对象和推导出来的循环边界产生对象。它表示一个初始的循环嵌套结构。C++层中的为所有语句（Statement）的容器。它的子类有，，，，，，，，，等等。该函数会处理schedule的依赖，核心部分是逆向遍历当中的（对于上面例子中就是先Compute Op，再两个Placeholder Op）。对于每个stage（除外），根据其attach type调用相应的逻辑。
   1. 对于上面的例子，Compute Op没有attach在其它计算中，因此它对应Stage的attach type为，因此这里调用函数产生。这步比较关键比较复杂，后面再展开。
   2. 然后通过对象（继承自）对前面生成的进行后处理。
4. 函数用于绑定buffer。它给每个参数张量分配buffer。如对于上面例子中的A, B, C三个张量，分别通过创建buffer并将之与张量绑定。
5. 函数基于产生的创建对象，它可以被用于TIR优化。代表包含了TIR statement的primitive function，它是low-level的代码表示。
6. 创建对象。基于上面生成的对象封装成对象并返回。一个可以有多个函数，比较简单的情况下就一个。

上面第函数中会调用函数针对ComputeOp对应Stage，返回一条由组成的pipeline，其大体流程相关代码如下：

函数主要步骤如下：

1. 函数主要创建对应的循环嵌套对应的那些对象并串成pipeline。
   1. 首先用函数检测计算类型。它遍历当前的所有当前有效对象，并根据它们的属性判断计算类型，对于上面的简单例子这里为。

      来源：ariesjzj

      声明：本站部分文章及图片转载于互联网，内容版权归原作者所有，如本站任何资料有侵权请您尽早请联系jinwei@zod.com.cn进行处理,非常感谢！