TVM学习（五）schedule

作者：安平博，Xilinx高级工程师；来源：AI加速微信公众号 Schedule是和硬件体系结构相关的一些列优化，Halide在其文章中对其做了以下定义： 1 When and where should be the value at each coordinate in each funcTIon be computed? 2 Where should they be stored? 3 How long are values cached and communicated across mulTIple consumers, and when are they independently recomputed by each? 第一条是描述了数据计算顺序对性能的影响，第二条是数据的存储位置对性能影响，最后一条是多线程处理过程中，不同线程数据应该如何进行交互。 参考文章：https://zhuanlan.zhihu.com/p/94846767，常用的shcedule有： 1 cache_read 将数据存储到片上缓存，减少访问数据时间。 2 cache_write 将结果写入片上缓存，然后再写入片外缓存。当然这里的片上和片外并不是绝对的概念，也可以理解为不同层次的存储结构。 3 set_scope 为数据指定存储位置，相比于cache_read和cache_write提供了更灵活的指定数据存储方式。本质上是相同的。 4 storage_align 在我看的文章中，storage_align是针对GPU shared memory的一个优化，目的是为了减少同一个bank的访问冲突。在GPU中shared memory被分割成多个bank，这些bank可以被独立线程同时访问。Storage_align就是为了将数据和bank大小匹配，减少bank conflict的发生。AI芯片中也有类似的问题，只有尽量减少bank冲突的发生，才能最大化并行计算。 5 compute_at 不懂CUDA，所以对文章中的代码不是很理解，但是从其解释看，对于多次循环的计算（或者多维计算），可以通过并行计算来降维。 6 compute_inline 将独立操作转化为内联函数，有点类似FPGA上的流水线计算。转化成内联函数从上层层面减少了stage。在FPGA中也有类似问题，可以将具有相同迭代的多条指令放在一起执行。 7 compute_root Compute_at的反操作。 8 fuse 将多个循环iter融合为一个iter。 9 split Fuse的反操作，将一次循环迭代拆分为多次。 10 reorder 调整循环计算迭代顺序。 11 TIle TIle也是将循环迭代进行拆分，拆分多次计算。是split+reorder。 12 unroll 将循环展开，增加并发执行。 13 vectorize 将循环迭代替换成ramp，可以通过SIMD指令实现数据批量计算，也就是单指令多数据计算。这在AI加速中会很常用，每条指令都是多数据计算的。 14 bind CUDA中使用的优化方法，将iter绑定到不同线程，实现并发计算。 15 parallel 实现多设备并行. 16 pragma 可以在代码中人为添加编译注释，人为干预编译优化。HLS中就是通过这样的方式来实现c的硬件编程的。 17 prefetch 将数据计算和load后者store数据重叠起来，在FPGA中是很常见优化方法。 18 tensorize 将tensor作为一个整体匹配硬件的计算核心，比如一个卷积运算就可以实现在FPGA上的一个匹配。 文章https://zhuanlan.zhihu.com/p/166551011 是通过官网的一个例子来介绍schedule的。在这个例子中，首先利用te的节点表达式建立了计算函数，然后调用create_schedule来创建schedule实例，然后再调用lower函数实现schedule优化。代码如下： # declare a matrix element-wise multiply A = te.placeholder((m, n), nam) B = te.placeholder((m, n), nam) C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], nam) s = te.create_schedule([C.op]) # lower will transform the computation from definition to the real # callable function. With argument `simple_mode=True`, it will # return you a readable C like statement, we use it here to print the # schedule result. print(tvm.lower(s, [A, B, C], simple_mode=True)) 我这里依然延续上一章的内容，看代码中关于schedule的处理。 在上一章我们在codegen生成中，通过以下调用链转到了schedule的处理。Codegen -> VisitExpr(CallNode* op) -> relay.backend._CompileEngineLower -> LowerInternal。LowerInternal函数为： 如果是外部定义的编译器，就只是建立cache_node节点和cache_func。如果是使用内部编译器，就会调用CreateSchedule建立schedule。接下来调用链为CreateSchedule -> ScheduleGetter.create -> te::create_schedule -> Schedule。create_schedule函数调用在文件re/schedule.h和te/schedule_lang.cc中。 create_schedule中主要有两件工作： 1 创建ReadGraph，获取post-dfs顺序的算符图。 2 初始化stage。 TVM中引入了stage的概念，一个op相当于一个stage，schedule优化是对stage的一个更改，可以增加，删减，更改其特性等。 通过createReadGraph可以遍历op图，返回op和其依赖的tensor列表。和遍历有关的主要函数为： Op -> InputTensors -> PostOrderVisit -> IRApplyVisit，在IRApplyVisit中定义了VisitExpr和VisitStmt函数用于遍历节点。 Stmt节点通常是节点中的主体实现，PrimExpr是TIR中节点的一个简单表达式。比如if节点： ReadGraph创建完成后，通过PostDFSOrder来获取post-dfs列表，其函数具体实现在graph.cc中， 通过不断迭代来进行深度优先搜索。 接下来是对stage进行初始化。 首先对postorder中的所有op初始化一个stage对象。我们看以下stage的定义： Stage类中主要定义了set_scope, compute_at, compute_root, bind, split, fuse等几种优化算法。同时定义了StageNode，在StageNode中定义了和优化相关的变量，包括op，iter变量等。看一下stage初始化代码： 关键的几个变量lef_iter_vars，all_iter_vars，这些有什么作用还需要深入看优化函数的代码。我们看几个schedule函数，先看一个最简单的：compute_inline。代码只有一行： (*this)->attach_type = kInline 对于标记了kInline的节点，在lower的时候会进行处理。应该会将其直接和调用的节点结合，合并两个op。 再看fuse函数，其代码为： IterVar表示计算中坐标轴，比如一个两级循环，每级循环就是一个axis。从代码中看出，fuse函数会对输入的所有axis进行合并，用fused变量替换合并后的axis。 这块代码比较抽象，先熟悉以下流程，之后再深入读一下。