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), name="A")

B = te.placeholder((m, n), name="B")

C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name="C")

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。

这块代码比较抽象,先熟悉以下流程,之后再深入读一下。

最新文章

最新文章