TVM中schedule介紹

作者：安平博，Xilinx高級工程師；來源：AI加速微信公眾號

Schedule是和硬件體系結(jié)構(gòu)相關(guān)的一些列優(yōu)化，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?

第一條是描述了數(shù)據(jù)計(jì)算順序?qū)π阅艿挠绊?，第二條是數(shù)據(jù)的存儲位置對性能影響，最后一條是多線程處理過程中，不同線程數(shù)據(jù)應(yīng)該如何進(jìn)行交互。

參考文章：https://zhuanlan.zhihu.com/p/94846767，常用的shcedule有：

1 cache_read

將數(shù)據(jù)存儲到片上緩存，減少訪問數(shù)據(jù)時(shí)間。

2 cache_write

將結(jié)果寫入片上緩存，然后再寫入片外緩存。當(dāng)然這里的片上和片外并不是絕對的概念，也可以理解為不同層次的存儲結(jié)構(gòu)。

3 set_scope

為數(shù)據(jù)指定存儲位置，相比于cache_read和cache_write提供了更靈活的指定數(shù)據(jù)存儲方式。本質(zhì)上是相同的。

4 storage_align

在我看的文章中，storage_align是針對GPU shared memory的一個(gè)優(yōu)化，目的是為了減少同一個(gè)bank的訪問沖突。在GPU中shared memory被分割成多個(gè)bank，這些bank可以被獨(dú)立線程同時(shí)訪問。Storage_align就是為了將數(shù)據(jù)和bank大小匹配，減少bank conflict的發(fā)生。AI芯片中也有類似的問題，只有盡量減少bank沖突的發(fā)生，才能最大化并行計(jì)算。

5 compute_at

不懂CUDA，所以對文章中的代碼不是很理解，但是從其解釋看，對于多次循環(huán)的計(jì)算（或者多維計(jì)算），可以通過并行計(jì)算來降維。

6 compute_inline

將獨(dú)立操作轉(zhuǎn)化為內(nèi)聯(lián)函數(shù)，有點(diǎn)類似FPGA上的流水線計(jì)算。轉(zhuǎn)化成內(nèi)聯(lián)函數(shù)從上層層面減少了stage。在FPGA中也有類似問題，可以將具有相同迭代的多條指令放在一起執(zhí)行。

7 compute_root

Compute_at的反操作。

8 fuse

將多個(gè)循環(huán)iter融合為一個(gè)iter。

9 split

Fuse的反操作，將一次循環(huán)迭代拆分為多次。

10 reorder

調(diào)整循環(huán)計(jì)算迭代順序。

11 tile

Tile也是將循環(huán)迭代進(jìn)行拆分，拆分多次計(jì)算。是split+reorder。

12 unroll

將循環(huán)展開，增加并發(fā)執(zhí)行。

13 vectorize

將循環(huán)迭代替換成ramp，可以通過SIMD指令實(shí)現(xiàn)數(shù)據(jù)批量計(jì)算，也就是單指令多數(shù)據(jù)計(jì)算。這在AI加速中會很常用，每條指令都是多數(shù)據(jù)計(jì)算的。

14 bind

CUDA中使用的優(yōu)化方法，將iter綁定到不同線程，實(shí)現(xiàn)并發(fā)計(jì)算。

15 parallel

實(shí)現(xiàn)多設(shè)備并行.

16 pragma

可以在代碼中人為添加編譯注釋，人為干預(yù)編譯優(yōu)化。HLS中就是通過這樣的方式來實(shí)現(xiàn)c的硬件編程的。

17 prefetch

將數(shù)據(jù)計(jì)算和load后者store數(shù)據(jù)重疊起來，在FPGA中是很常見優(yōu)化方法。

18 tensorize

將tensor作為一個(gè)整體匹配硬件的計(jì)算核心，比如一個(gè)卷積運(yùn)算就可以實(shí)現(xiàn)在FPGA上的一個(gè)匹配。

文章https://zhuanlan.zhihu.com/p/166551011 是通過官網(wǎng)的一個(gè)例子來介紹schedule的。在這個(gè)例子中，首先利用te的節(jié)點(diǎn)表達(dá)式建立了計(jì)算函數(shù)，然后調(diào)用create_schedule來創(chuàng)建schedule實(shí)例，然后再調(diào)用lower函數(shù)實(shí)現(xiàn)schedule優(yōu)化。代碼如下：

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

我這里依然延續(xù)上一章的內(nèi)容，看代碼中關(guān)于schedule的處理。

在上一章我們在codegen生成中，通過以下調(diào)用鏈轉(zhuǎn)到了schedule的處理。Codegen -> VisitExpr(CallNode* op) -> relay.backend._CompileEngineLower -> LowerInternal。LowerInternal函數(shù)為：

如果是外部定義的編譯器，就只是建立cache_node節(jié)點(diǎn)和cache_func。如果是使用內(nèi)部編譯器，就會調(diào)用CreateSchedule建立schedule。接下來調(diào)用鏈為CreateSchedule -> ScheduleGetter.create -> te::create_schedule -> Schedule。create_schedule函數(shù)調(diào)用在文件re/schedule.h和te/schedule_lang.cc中。

create_schedule中主要有兩件工作：

1 創(chuàng)建ReadGraph，獲取post-dfs順序的算符圖。

2 初始化stage。

TVM中引入了stage的概念，一個(gè)op相當(dāng)于一個(gè)stage，schedule優(yōu)化是對stage的一個(gè)更改，可以增加，刪減，更改其特性等。

通過createReadGraph可以遍歷op圖，返回op和其依賴的tensor列表。和遍歷有關(guān)的主要函數(shù)為：

Op -> InputTensors -> PostOrderVisit -> IRApplyVisit，在IRApplyVisit中定義了VisitExpr和VisitStmt函數(shù)用于遍歷節(jié)點(diǎn)。

Stmt節(jié)點(diǎn)通常是節(jié)點(diǎn)中的主體實(shí)現(xiàn)，PrimExpr是TIR中節(jié)點(diǎn)的一個(gè)簡單表達(dá)式。比如if節(jié)點(diǎn)：

ReadGraph創(chuàng)建完成后，通過PostDFSOrder來獲取post-dfs列表，其函數(shù)具體實(shí)現(xiàn)在graph.cc中，

通過不斷迭代來進(jìn)行深度優(yōu)先搜索。

接下來是對stage進(jìn)行初始化。

首先對postorder中的所有op初始化一個(gè)stage對象。我們看以下stage的定義：

Stage類中主要定義了set_scope, compute_at, compute_root, bind, split, fuse等幾種優(yōu)化算法。同時(shí)定義了StageNode，在StageNode中定義了和優(yōu)化相關(guān)的變量，包括op，iter變量等?？匆幌聅tage初始化代碼：

關(guān)鍵的幾個(gè)變量lef_iter_vars，all_iter_vars，這些有什么作用還需要深入看優(yōu)化函數(shù)的代碼。我們看幾個(gè)schedule函數(shù)，先看一個(gè)最簡單的：compute_inline。代碼只有一行：

(*this)->attach_type = kInline

對于標(biāo)記了kInline的節(jié)點(diǎn)，在lower的時(shí)候會進(jìn)行處理。應(yīng)該會將其直接和調(diào)用的節(jié)點(diǎn)結(jié)合，合并兩個(gè)op。

再看fuse函數(shù)，其代碼為：

IterVar表示計(jì)算中坐標(biāo)軸，比如一個(gè)兩級循環(huán)，每級循環(huán)就是一個(gè)axis。從代碼中看出，fuse函數(shù)會對輸入的所有axis進(jìn)行合并，用fused變量替換合并后的axis。

這塊代碼比較抽象，先熟悉以下流程，之后再深入讀一下。

審核編輯：何安