# Machine Learning Compilation_第三週_張量程序抽象案例研究：TensorIR(0702)
###### tags: `陳天奇` `MLC` `Machine Learning Compilation`
[官網連結](https://mlc.ai/summer22-zh/)
[課程連結](https://mlc.ai/summer22-zh/schedule)
[課程Colab](https://github.com/mlc-ai/notebooks/blob/main/3_TensorIR_Tensor_Program_Abstraction_Case_Study_Action.ipynb)
[課程練習](https://mlc.ai/zh/chapter_tensor_program/tensorir_exercises.html)
課程說明：這門課程並不要求要有tvm的相關基礎，只需要你有一些python、機器學習或是深度學習框架的基礎就可以，課程主要是希望可以帶大家從零到有的學習tvm。也因為上一堂課的反饋，所以這一週的課程會就tvm來聊聊。
另外，這節課基本上是用colab直接線上demo，沒有簡報，可以直接開啟[課程Colab](https://github.com/mlc-ai/notebooks/blob/main/3_TensorIR_Tensor_Program_Abstraction_Case_Study_Action.ipynb)。
前面的課程中我們已經知道，機器學習編譯是一個張量程序的變換，這個張量程序可以是一個算子，也可以是多個算子，而我們所關注的部份在於，我們可以怎麼樣的設計一個抽象來表述想要優化的程序。課程中也會利用範例來說明如何設計一個抽象。
首先，如果你有打開Colab就會看到，範例中有兩個矩陣，兩個矩陣在做完相乘之後會經過ReLU，這是一個很標準的線性計算之後接非線性計算的範例：
* $Y_{ij} = \sum_k A_{ik} B_{kj}$
* $C_{ij} = \text{relu}(Y_{ij}) = \max(Y_{ij}, 0)$
相同的作法如果在numpy上做的話，會是下面這樣：
```python
dtype = 'float32'
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
# a_np @ b_np等價於np.matmul(a_np, b_np)
c_mm_relu = np.maximum(a_np @ b_np, 0)
```
雖然我們用numpy只是簡單的一行就可以處理，不過底層其實numpy幫我們處理掉很多麻煩事。平常我們是不怎麼關心到底我們呼叫`np.maximum`的時候底層到底numpy做了什麼。不過在機器學習編譯的時候，因為我們想要優化某些過程，所以我們就會需要知道底層做了什麼。
範例的選擇也是有意義的，普遍我們常見的作法就是先一層線性，再一層ReLU。所以課程中我們就要來做一個挑戰，一樣是用python，一樣是用numpy，但是我們不用高度封裝好的那些函數，而是使用低階的指令來實現。也因此，可以看的到課程所提供的colab中有一個`lnumpy_mm_relu`的函數。這麼做的好處在於，能寫到這麼底層的東西，轉到C就不是太大的問題：
```python
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((128, 128), dtype='float32')
for i in range(128):
for j in range(128):
for k in range(128):
if k == 0:
Y[i, j] = 0
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
for i in range(128):
for j in range(128):
C[i, j] = max(Y[i, j], 0)
```
函數中有三個參數，`A, B`為輸入，`C`為輸出。`Y`是結果的暫存，也就是`A, B`兩個矩陣做相乘之後的結果。這部份可以參考[維基百科](https://zh.wikipedia.org/wiki/%E7%9F%A9%E9%99%A3%E4%B9%98%E6%B3%95)的圖解說明會比較好理解：
接下來的一段程式主要是拿來做為驗證使用，驗證我們的計算結果是跟numpy的函數計算結果是一致的(講的真的夠詳細的)：
```python
c_np = np.empty((128, 128), dtype=dtype)
# 執行之後c_np就會有計算後的結果
lnumpy_mm_relu(a_np, b_np, c_np)
# numpy內建的測試函數
np.testing.assert_aliclose(c_mm_relu, c_np, rtol=1e-5)
```
我們所設置的函數`lnumpy_mm_relu`已經包含很多上一節課講到的基本元素在裡面，像是loop，這就給了我們優化的空間。所以接下來就可以開始處理上一節課所談到的TensorIR：
```python
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np
@tvm.script.ir_module
class MyModule:
@T.prim_func
def mm_relu(A: T.Buffer[(128, 128), 'float32'],
B: T.Buffer[(128, 128), 'float32'],
C: T.Buffer[(128, 128), 'float32'],):
"""
@tvm.script.ir_module: 包裝成python對應的object
@T.prim_func: 宣告為一個元張量函數
T.Buffer：這主要用來保存臨時的內存數組，一樣會需要它的維度以及資料類型
T.alloc_buffer：在中間開一個臨時的數組
for i, j ,k = T.grid就對應numpy的三個迴圈
T.init：如果是第一次處理就視為初始化
with T.block("Y")是numpy裡面很明顯沒有的，這是TensorIR本身一個特殊的構造
主要用來宣告一個計算單元，並且存在一個特性，isolation，
這個範例來看就是讀取A、B來計算Y
T.axis：block本身的座標軸
T.axis.spatial: 對應矩陣中某個位置，是個實際存在的維度，以我們這個範例來說
就會是Y或C，所以vi, vj對應的就會是Y、C的實際維度
討論區來看有人說稱為space loop或是space axis
vi, vj的執行順序並不是那麼重要
T.axis.reduce: 並非實際存在的維度，可能只是過程中需要的，是會消失的，
所以又稱為reduction loop
而且你必需要確實的執行之後才能得到結果
T.func_attr({'global_symbol': 'mm_relu', 'tir.noalias': True})
1. global_symbol為該函數生成之後的名稱，
2. tir.noalias為宣告內部buffter指針不重覆
"""
T.func_attr({'global_symbol': 'mm_relu', 'tir.noalias': True})
Y = T.alloc_buffer((128, 128), dtype='float32')
for i, j, k in T.grid(128, 128, 128):
with T.block('Y'):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block('C'):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
```
兩個放一起比較如下圖：
幾點說明如下：
1. 兩邊都各有三個參數`A,B,C`
2. 都有設置一個`Y`做為暫存計算，也看的出來宣告的方式有些許的不一樣
3. numpy版本中的`i, j, k`迴圈在TensorIR中是以grid來表示
4. `Y[i, j]`的計算也是直接的對應
5. `if k==0: Y[i, j] = 0`即對應`with T.init(): Y[vi. vj] = T.float32(0)`
6. `with T.block("Y")`很明顯是numpy版本所沒有的，主要用來宣告一個計算單元，這個範例來看就是讀取A、B來計算Y，舉例來說，`Y[0, 1]`的值就會是`vi=0, vj=1, vk=0,1,2,3,4,5,...,127`，整個計算之後就可以得到`Y[0, 1]`的結果，基本上`vi, vj`的執行順序是沒有關係的，這從上圖的矩陣計算不難聯想，這就很容易可以達到平行化計算
後續會說明一些自動化分析的方式來取得相關的程式轉換，不過這只適用於一些簡單的情況。`T.block`有著一種隔離(isloation)的特性，也就是說，我們就是在這個`block`中宣告要做那些事，這些事也就只會在這個`block`中處理。理論上從這個`block`中的訊息我們就可以猜的出外部循環要做的事情是什麼。要注意的是，如果`block`的維度設置與外部循環的設置不一的時候(舉例來說，`T.axis.spatial(128)`設置為`127`)那就會報錯。`block`本身即為TensorIR設計中的核心。
如果說，每次的宣告我們都要仔細的一行一行去定義`T.axis.spatial`那會有點累：
```python
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
```
所以我們可以採用一行流來取代這個作法：
```python
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
```
```python
@tvm.script.ir_module
class MyModule:
@T.prim_func
def mm_relu(A: T.Buffer[(128, 128), 'float32'],
B: T.Buffer[(128, 128), 'float32'],
C: T.Buffer[(128, 128), 'float32'],):
"""
@tvm.script.ir_module: 包裝成python對應的object
@T.prim_func: 宣告為一個元張量函數
T.Buffer：這主要用來保存臨時的內存數組，一樣會需要它的維度以及資料類型
T.alloc_buffer：在中間開一個臨時的數組
for i, j ,k = T.grid就對應numpy的三個迴圈
T.init：如果是第一次處理就視為初始化
with T.block("Y")是numpy裡面很明顯沒有的，這是TensorIR本身一個特殊的構造
主要用來宣告一個計算單元，並且存在一個特性，isolation，
這個範例來看就是讀取A、B來計算Y
T.axis：block本身的座標軸
T.axis.spatial: 對應矩陣中某個位置，是個實際存在的維度，以我們這個範例來說
就會是Y或C，所以vi, vj對應的就會是Y、C的實際維度
討論區來看有人說稱為space loop或是space axis
vi, vj的執行順序並不是那麼重要
T.axis.reduce: 並非實際存在的維度，可能只是過程中需要的，是會消失的，
所以又稱為reduction loop
而且你必需要確實的執行之後才能得到結果
T.func_attr({'global_symbol': 'mm_relu', 'tir.noalias': True})
1. global_symbol為該函數生成之後的名稱，
2. tir.noalias為宣告內部buffter指針不重覆
"""
T.func_attr({'global_symbol': 'mm_relu', 'tir.noalias': True})
Y = T.alloc_buffer((128, 128), dtype='float32')
for i, j, k in T.grid(128, 128, 128):
with T.block('Y'):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block('C'):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
```
如果我們嚐試著去查看這個類別的類型的話，它會是一個`tvm.ir.module.IRModule`，這也是TensorIR的核心之一。課程中我們提過，機器學習編譯是一種張量函數轉換的過程，如何用一個python類別來表示，靠的就是這個`IRModule`。
在複雜一點的情況下，一個IRModule會包含兩個張量函數，以我們的課程中的範例來看就包含矩陣乘法與ReLU。到這邊已經基本說明TensorIR的基本概念。
下半場的開始就先給出一個不同版本的矩陣相乘之後接ReLU的numpy範例：
```python
def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((128, 128), dtype='float32')
for i in range(128):
for j0 in range(32):
for k in range(128):
for j1 in range(4):
j = j0 * 4 + j1
if k == 0:
Y[i, j] = 0
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
for i in range(128):
for j in range(128):
C[i, j] = max(Y[i, j], 0)
c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu_v2(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)
```
很明顯的，跟最一開始的版本最大的差異就是把`j`的loop拆成`j0, j1`兩個，也因此`j = j0 * 4 + j1`：
上圖說明的是`j = j0 * 4 + j1`。
:::info
課程的函數函數code2html使用[pygments](`https://pygments.org/`)來做為程式碼關鍵字highlight的處理。
:::
接下來我們就要針對v2版本的部份做TensorIR的處理：
1. 首先建立一個`Schedule`類別
2. 從得到的`Schedule`類別中，利用`get_block`從中取得要處理的那個block
3. 再從這個拿到的block中利用`get_loops`取得裡面的相關迴圈
4. 利用`split`拆解迴圈，將`j`拆解為兩個迴圈，`j0, j1`
5. 最後再利用`reorder`來調整迴圈的順序
處理過程中我們可以利用`mod.script()`來確認當下的變化，其實可以發現，調整之後剩下的部份都會是由框架幫你處理掉。基本上如果你跟著jyputer notebook的處理來看，最終得到的會等價於我們用numpy硬寫一發的v2版本。
課程的最後還有一個最終版本，是將`C`的部份放到計算過程中，而不是最後再用一個迴圈來處理，少了一個迴圈就可以減少不少計算效能。有時候想想，迴圈真的是必要之惡就是了，不過依稀記得演算法老師上課說過，用遞迴思考，用迴圈實作。
在`Getting to Another Variant`的部份可以看到：
```python=
block_C = sch.get_block('C', 'mm_relu')
sch.reverse_compute_at(block_C, j0)
sch.mod.show()
```
第2行的部份意謂著要將`block_C`的計算拉到上面`j0`去：
這類的操作即為『算子融合』，記得到課程提供的ipynb去實際點擊執行確認變化：
有融合也有拆解，拆解的部份可以利用`decompose_reduction`來實作：
```python=
block_Y = sch.get_block('Y', 'mm_relu')
sch.decompose_reduction(block_Y, k)
sch.mod.show()
```
拆解之後如下：
最終結果所對應的numpy的話就會是：
```python=
def lnumpy_mm_relu_v3(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((128, 128), dtype="float32")
for i in range(128):
for j0 in range(32):
# Y_init
for j1 in range(4):
j = j0 * 4 + j1
Y[i, j] = 0
# Y_update
for k in range(128):
for j1 in range(4):
j = j0 * 4 + j1
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
# C
for j1 in range(4):
j = j0 * 4 + j1
C[i, j] = max(Y[i, j], 0)
c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu_v3(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)
```
其實老師真的很有心，一邊寫tensorIR的實作，還一邊提供numpy自己硬刻一發的對應。不過這兩邊的差別在於，如果是利用tensorIR來操作的話，我們可以利用調整某一個block的方式來做，但numpy的話你可能要整個重刻。
課程最後說明，在一切就緒之後，可以利用`tvm.build`在給定模組與目標之後建置執行需求：
```python=
# MyModule是最原始的版本
rt_lib = tvm.build(MyModule, target="llvm")
```
`tvm`本身也有相對應的`nd.array`類型，可以直接以numpy做為input來轉換，以此定義input、output：
```python=
# input
a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
# output
c_nd = tvm.nd.empty((128, 128), dtype='float32')
```
然後就可以先利用`rt_lib`取得要執行的目標函數，並執行：
```python=
# 取得執行函數
func_mm_relu = rt_lib('mm_relu')
# 執行
func_mm_relu(a_nd, b_nd, c_nd)
```
當然你也可以拿課程中最後調整的版本來做轉換：
```python=
rt_lib_after = tvm.build(sch.mod, target='llve')
```
你還可以用內建函數來比較調整前後的執行效率，這個tvm的內建函數會去掉第一次的執行，然後再取平均：
```python=
f_timer_before = rt_lib.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of MyModule %g sec" % f_timer_before(a_nd, b_nd, c_nd).mean)
f_timer_after = rt_lib_after.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of transformed sch.mod %g sec" % f_timer_after(a_nd, b_nd, c_nd).mean)
```
這兩者之間的差異可以從課程中提供的ipynb中看的到，主要在於不同的迭代順序造成不同的執行效率：
每次迭代中我們會小批小批的去取資料，這個小批的取資料動作就會造成效率上的差異，這跟當代的CPU架構是有關的：
上圖可以看的到，資料在DRAM跟在L2-cache、L1-cache的執行效率有極大的差異，所以CPU在執行上會做一個假設，它會假設當你取某一個記憶體區塊的資料的時候，你也會需要這附近區域的資料，所以它會將這附近區域的資料就這樣的丟到L1-cache。這意謂著，如果你在取DRAM資料是連續存取的話，那你去從L1-cache取到資料的機率就會高很多，效率自然就快。
也因此我們如果把程式碼的迭代拉出來看的話是可以發現到，迭代過程中一直在看的就是j1裡面的元素，而這個元素是已經被拉到L1-cache，執行起來自然就快：
相同的道理也可以在對`C`計算的調整中看的到，我們最後面是把`C`的計算拉進去迭代中處理，而不是在整個結果結束之後再又重啟一次迭代，因為資料已經在L1-cache，執行起來就是快。
課程的最後老師提供一個練習，讓我們去測試一下不同`j`大小執行效率上的差異，想法上應該是不同的CPU就會有不同的差異。
最後的最後，tvm本身還提供張量表達式讓我們設計張量流：
```python=
from tvm import te
A = te.placeholder((128, 128), "float32", name="A")
B = te.placeholder((128, 128), "float32", name="B")
k = te.reduce_axis((0, 128), "k")
Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")
```
其中`Y`計算的過程即為$Y_{ij} = \sum_k A_{ik} B_{kj}$，然後再把這個表達式利用內建函數做一個包裝轉換，你仍然可以得到我們最原始的那個`MyModule`的設計：
```python=
te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
MyModuleFromTE = tvm.IRModule({"mm_relu": te_func})
MyModuleFromTE.show()
```
結果如下：
```python=
@tvm.script.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer[(128, 128), "float32"], B: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"]) -> None:
# function attr dict
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
# body
# with T.block("root")
Y = T.alloc_buffer([128, 128], dtype="float32")
for i0, i1, i2 in T.grid(128, 128, 128):
with T.block("Y"):
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
T.reads(A[i, k], B[k, j])
T.writes(Y[i, j])
with T.init():
Y[i, j] = T.float32(0)
Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
for i0, i1 in T.grid(128, 128):
with T.block("C"):
i, j = T.axis.remap("SS", [i0, i1])
T.reads(Y[i, j])
T.writes(C[i, j])
C[i, j] = T.max(Y[i, j], T.float32(0))
```
總結這禮拜的課程，過去的開發模式我們可能就是不斷做著重覆開發的事情，version 1.0、1.1、...之類的。但是在機器學習編譯中我們強調的是Transform(轉換)，可能我們寫出一個初始版本之後(課程中的MyModule)，透過轉換取得IRModule，接下來我們做的就是單純的對這個IRModule中的張量函數做一系列的轉換，迭代的block的轉換，算子函數的融合也是一種轉換，轉換之後得到一個新的版本，然後再去執行`build`，以取得我們預計佈署目標的形式：
還是那句話，老師真的很有心，說的這麼細，又給出這麼漂亮的範例，真心感謝。
Sign in with Wallet
Connect another wallet