# Machine Learning Compilation_第二週(0625)
###### tags: `陳天奇` `MLC` `Machine Learning Compilation`

[官網連結](https://mlc.ai/summer22-zh/)
[課程連結](https://mlc.ai/summer22-zh/schedule)
[TVM深度學習編譯技術_工研院董明智_清大劉靖家教](https://ictjournal.itri.org.tw/content/Messagess/contents.aspx?&MmmID=654304432061644411&CatID=654313611227352166&MSID=1073702176133520045)

本週課程談的是MLC比較核心，也是比較基礎的架構，稱之為張量算子顯示
## Recap: Key Elements in Machine Learning Compilation

上次的課程如果還有印象的話有提過MLC的關鍵元素，也就是tensor與tensor function。也談到過機器學習編譯(MLC)處理過程就是在tensor function不同的抽象上的變換(可見上週課程筆記說明)。

## Primitive Tensor Function

原張量函數，或者可以稱為張量算子。上圖左是一個張量的計算流程，先是一個線性計算，然後加上bias，經過relu再做一次線性計算，最後是softmax輸出。

所以我們知道，機器學習的計算過程是分很多步的，每一步的單元步驟(像是linear、add)就對應於Primitive Tensor Function。

## Primitive Tensor Function in ML Frameworks

這種Primitive Tensor Function在現行的很多框架中都可以直接呼叫函數來操作，上圖為例，使用的是pytorch，定義兩個tensor(a、b)，然後直接使用`torch.add(a, b, out=c)`來計算結果。基本上這個`add`就可以看成是一個基本的操作單元，也就是Primitive Tensor Function。

問題來了，那我們應該如何去實現這個torch（？這句不是很清楚）？接下來就是，實際佈署上是會有很多步，不會只有一個`add`，那我們應該如何把這些步驟合在一起？

今天就用這一個`add`來探討怎麼做比較好！

## Abstractions for Primitive Tensor Function

就跟上週課程有提到的兩個部份，那就是抽象與實現。對於同一個tensor function來說，我們有很多種方式可以實現。就像上圖最左邊的`torch.add`，這也是一種抽象的表示，非常高階的抽象，我們只要呼叫它就可以執行`add`的動作。我們可以把這個表示再更低階的抽象，像是中間圖用python所寫出來的函式。當然實務上可能會用c語言，那就可以用右邊那個c語言的實現。我們可以發現到，就這麼一個`add`就可以有很多種不同的表示。

## MLC via Primitive Function Transformation

在機器學習編譯的過程中最常見的一種優化就是對單元算子做一個替換。就像上圖的範例所呈現那般，左邊的`add`是原本以python的一個實現，也許我們可以用右邊的方式，平行化計算的概念來實現(這只是一個概念說明，右邊的偽代碼並不代表任何的程式語言)，也許底層的硬體支援長度為四個向量計算，那我們就可以一次做四個元素的計算。

以這個範例來看，它的接口本身並沒有任何的變換，所以我們要考慮的就只有怎麼對單元算子做變換。

## Primitive Function Transformation

對於算子變換與優化本身的話有很多種做法，常見作法是將其直接映射到相對應的算子函數庫上。就像簡報提到的範例，如果是一個`add`要丟到cuda的話，那就可以直接丟到`cudaAdd`就可以。如果是AMD或是ARM的CPU，那就映射到相對應的算子函數庫去。

另外一種方式就是做更為細緻的程式變換。不同的算子變換對於抽象本身就有不同的要求。課程中我們就會來談到這一點。

那為什麼我們會需要做這樣的算子換變呢？主要是我們希望可以找到一些等價的城市可以在我們的目標應用環境上有更好的表現。

## Tensor Program Abstraction

剛剛提了單元算子，這邊要談的是張量程序抽象，也就是Tensor Program Abstraction。

上一張簡報提到，如果我們只是單純的做庫函數的映射的話，那我們就只是去呼叫一個函數而以，並不需要把整個單元算子展開。不過如果想要對單元算身做進一步的變化，或是優化的話，那就涉及另一種抽象，我們就稱為張量函數抽象。

上圖給出的是一個`tvm`。

> [tvm介紹](https://ictjournal.itri.org.tw/content/Messagess/contents.aspx?&MmmID=654304432061644411&CatID=654313611227352166&MSID=1073702176133520045)

## Key Elements of a Tensor Program

張量函數程序中會有三個關鍵元素：(1)多維的輸入、輸出的數組、(2)伴隨著一些迴圈、(3)對應迴圈會有一些計算。

大致上所有的張量函數程序都會有這三個元素，而它的循環、計算本身有一些特殊性，以上圖的迴圈為例，每個計算都是獨立的，不管從前往後迴圈、還是從後往前迴圈都是一樣的結果，也因為這樣的特性，我們可以把張量算子與張量程序的抽象拿到外部做一些等價的變換。

不過為什麼我們不直接寫C語言或是cuda呢？整個機器學習編譯(MLC)不光是人去寫，我們還要能夠有過程可以變換，也就是給定一個張量算組，由原始的形態變換到一個相對比較優化的形態，也因此我們需要一個程序來協助我們做這種等價的變換。

## Why do we need Tensor Program Abstraction

舉個範例來說，上圖左是原始的程序，右邊的話就是我們覺得比較好的轉換後的程序，因為它有著平行計算，計算過程中也以向量化來處理，每次以長度為四做為單位來處理。

當然，我們可以直接自己怒刻一發來處理右邊這個轉換後的程序，只是如果能夠有辦法直接做一個等價的轉換就直接達成，那是不是更省事？可以看的到，Program-based transformations下面有一個轉換的操作，通過這個簡單的操作，我們就可以拿到等價右邊轉換後的結果。

## Example Transformation: Loop Splitting

這邊給出我們能做的轉換操作。

左邊的Code是原始的程式，右邊的Transformation是我們做的轉換操作。主要是將一個loop切成兩個，也就是Code下方所指出的程式。當然，如果只是單純的做loop的拆解的話，對於效能本身是不大有影響的。只是這是常見的第一步。

## Example Transformation Loops: Loop Reorder

拆解之後的loop我們就可以做reorder，也就是重新排序。要注意，這時候的xi、xo的順序已經變成，原本是先`range(32)`，現在調整之後已經是先`range(4)`。

## Example Transformation: Thread Binding

這時候我們就可以去做gpu的thread binding。後面課程會再詳細說明gpu上的計算是怎麼一回事。

問題來了，為什麼我們要做這些轉換？轉來轉去做的還是原本的加法，只是實現的過程有點不一樣而以。也就是這個不一樣，取決於底層的硬體，每一種做法在不同的硬體上可能就有不同的執行效率。

## We cannot Arbitrarily Transform any Program

當然，並不是說因為上面說那麼多，我們就可以隨隨便便來搞兩下的。以上圖為例，這個範例跟上一個很像，但最大差異在於，後面的$B$的計算裡面有用到上一個迭代循環的元素來計算($xi - 1$)，這種情況下迴圈的順利就很有關係了。某種程度來說，也許這種情況我們也可以在一開始設計張量程序表示的時候就把這些因素給考慮進去，然後嚐試引入一些額外的訊息。

## Extra Structure in Tensor Program Abstraction

這邊說明的就是剛剛說的嚐試引入一些額外的結構，可以看到`vi=T.axis.spatial(128, i)`，這段語法說的是，`vi`對應一個迭代器，這個迭代器會從0迭代到128(上面的`for`迴圈)，然後`spatial`則表示這是可以平行化處理的，也就是不同維度上是相互獨立執行的。

## Tensor Program Abstraction in Action

這邊開始的課程是實際的執行畫面，主要說明如何構造對應的張量程序。裡面會用到的library為`tvm`，不過課程中也直接提供一個package，`python3 -m pip install mlc-ai-nightly -f https://mlc.ai/wheels`。課程內容會逐漸的往`tvm`去說明。


我們可以嚐試利用標量算子程序來構造一個對應的程序。上圖是剛剛已經有提過的範例，裡面用的是一種稱為『tvm script』的東西，然後我們嚐試在python使用它。

上面的程式中所構建的class，`MyModule`，我們用一個decorator `@tvb.script.ir_module`來修飾它，如果你去查這個class的`type`的話，它會是`tvm.ir.module.IRModule`，我們可以將之視為一個集合，一個包含張量算子程序的集合，或者說是張量程序的集合。

程式碼來看，`IRModule`這個集合很簡單，就只是包含`main`這個函數而以。如果我們想看這個`IRModule`的程式碼是怎麼一回事，我們可以透過method `script`來確認，以上面class為例的話，我們可以執行`MyModule.script()`：

如果想做程序的轉換，可以有幾個操作如下：

* 我們可以利用`sch = tvm.tir.Schedule(MyModule)`來處理，這可以讓我們做對應MyModule的操作
* 接下來可以執行`block_c = sch.get_block('C')`，這讓我們取得對應MyModule的`with T.block('C')`這個block
* 然後可以利用這個block再取得這個block外的迴圈，`i, = sch=get_loops(block_c)`，也就是`for i in range(128)`
* 現在我們可以嚐試來修正這個迴圈，利用`i0, i1, i2 = sch.split(i, factors=[None, 4, 4])`來做迴圈的拆分
* 相關的修正結果可以利用`print(sch.mod.script())`來確認，對於上面的修正結果比較的有感覺

上面的操作結束之後可以發現到，原本的迴圈是`for i in range(128)`，經過上面的調整已經變成是三塊了，我們在拆分的時候給的指令是`[None, 4, 4]`，因此最後面兩個區塊的部份程式會幫我們調整成我們指定的，而第一個維度的部份則是由程式處理，128 / 16 = 8。而`vi = T.axis.spatial(128, 1)`也變成`vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)`。

當然，我們說過迴圈的順序是可以調整，這部份可以利用`sch.reorder(i2, i1)`來做個驗證：

還有平行化的宣告，一樣可以利用`sch.parallel(i0)`來宣告，這意味著我們要把`i0`的部份平行化處理：

上圖的箭頭處可以看的到，`i_0`已經被定義為平行化處理，`for i_0 in tir.parallel(8)`。

這個過程說明著，我們建構一個MyModule，然後利用Schedule來包裝這個MyModule，再以此來做一系列的操作轉換，這樣的轉換就會造成不同的執行效率。


轉換之後的結果並不是一個可以直接執行的物件，我們還需要再根據要佈署的環境做一些編譯，這時候可以利用`rt_mod = tvm.build(sch.mod, target='llvm')`來處理，其中參數`target='llvm'`意謂著要在CPU上執行。原本的`sch.mod`只是一個用來做一系列轉換操作的模型，轉成`rt_mod`就是要執行用的(runtime)。然後再提出裡面的函數`main`，`func = rt_mod['main']`，這是一個把張量算子程序編譯之後取得的可執行函數(見上圖)。

這邊老師也給出一個很簡單的範例，根據函數需要的參數宣告三個變數：
* `a = tvm.nd.array(np.arange(128, dtype='float32'))`
* `b = tvm.nd.array(np.ones(128, dtype='float32'))`
* `c = tvm.nd.empty([128], dtype='float32')`

然後執行函數`func[a, b, c]`，給的參數記得應該上面class的定義，得到結果如下：

迴圈的拆分三份的部份可以自己手算一下，得到的結果是等價原始迴圈的執行，這是廢話，不過手算一下真的會比較有感覺。

## Summary

總結來看，今天講了兩個概念。首先，整個機器學習編譯的過程中，單元算子函數本身是一個非常重要的單位。我們對於機器學習編譯的優化就是把一個單元算子函數從一個表示轉換成另一種表示。最常處理到的就是迴圈的細膩度的變換，如上面課程中的實作所給出的範例。

另外就是，我們在做張量函數抽象的時候會需要加入一些額外的結構，加入額外結構的用意在於，我們希望可以保證轉換過程中的正確性，以及可以提供給我們更多信息。