作者 | OPEN AI LAB 研究員 呂春瑩
出品 | AI科技大本營
頭圖 | CSDN下載自視覺中國
隨着AI技術的快速發展，深度學習在各個領域得到了廣泛應用。深度學習模型能否成功在終端落地應用，滿足產品需求，一個關鍵的指標就是神經網絡模型的推理性能。於是，一大波算法工程師爲了算法的部署轉崗算子優化工程師。然而，優化代碼並不是一件簡單的事，它要求工程師既要精通計算機體系架構，又要熟悉算法的計算流程，於是，稍微有經驗的深度學習推理優化工程師都成了各家公司爭搶的“香餑餑”。相關人才少，但需求多，算子優化自動化成爲了未來的一大趨勢。
爲了方便更多的工程師進行推理優化，一個致力於降低優化門檻，提升優化開發效率的算子自動優化工具AutoKernel宣佈正式開源！
AutoKernel特色：
低門檻: 無需底層優化彙編的知識門檻
簡單易用: 提供docker環境，無需安裝環境，plugin一鍵集成到推理框架Tengine
高效率: 無需手寫優化彙編，一鍵生成優化代碼，一鍵部署
AutoKernel使用業界廣泛使用的自動代碼生成項目Halide，通過輸入計算描述和調度策略，自動生成底層代碼。AutoKernel支持以plugin的形式，將生成的自動優化算子一鍵部署到推理框架Tengine中。
下面，本教程將帶領大家一步步優化矩陣乘法GEMM。無需手工擼代碼，編寫繁雜冗長的底層彙編代碼，只需十幾行簡潔的調度代碼。
在詳細講解優化步驟前，我們先談談優化的本質。我們在談”優化“的時候，計算機底層做了什麼？優化的”瓶頸“是什麼？爲什麼通過一波”優化操作“，性能就能提升呢？AutoKernel使用的Halide是如何實現自動優化的呢？
要解答這些疑問，我們需要了解一下硬件的基礎的體系結構，瞭解硬件如何工作，才能在軟件上實現算法的時候，儘可能去考慮利用硬件的一些特性，來做到高效的、極致的優化。
上圖是典型的存儲理器層次結構：主存容量大，訪問速度慢，寄存器和緩存讀取速度快，但容量有限。在寄存器的層級上，CPU可以在一個時鐘週期內訪問它們，如果CPU去訪問外部的DDR的話，延遲是非常大的，大概是200個時鐘週期左右。如果CPU去訪問cache的話，一般需要6到12個cycle就夠了。所以，一個很重要的一個優化宗旨是：優化內存訪問，充分利用寄存器和高速緩存去存數據。
第二個優化宗旨則是提高並行性：充分利用SIMD進行指令向量化和多核心並行。大部分現代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數據流）。在同一個CPU循環中，SIMD可在多個值上同時執行相同的運算/指令。如果我們在4個數據點上進行向量化，一次計算四個數據，理論上就可以實現4倍的加速。
運行環境搭建
AutoKernel提供了docker鏡像，docker裏已經配置好運行環境，進入docker即可直接運行demo代碼：
# 拉取鏡像docker pull openailab/autokernel# 啓動容器，進入開發環境docker run -it openailab/autokernel /bin/bash# 獲取代碼git clone https://github.com/OAID/AutoKernel.gitcd AutoKernel/doc/tutorials/data/
目錄下的build.sh是demo的執行腳本，運行需要指定優化步驟step，可選的step是從1 到7，其中step= 1 是默認不優化的，step=7是最極致優化的。
優化效果
# 執行demo./build.sh 1./build.sh 7
下圖展示了在Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz的電腦上的優化效果，無需手工擼代碼，無需編寫繁雜冗長的底層彙編代碼，只需十幾行簡潔的調度代碼, 就能性能優化200+倍~
優化步驟
以下是更爲詳細的優化步驟：
STEP1
第一個步驟是不帶任何優化的。用Halide語言直接描述GEMM的計算過程。
Var x,y; RDom k(0, K); Func gemm("gemm"); gemm(x, y) += A(k, y) * B(x, k);
計算M=N=K=640的矩陣乘法。運行腳本第一個參數指定step=1。耗時結果如下：
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 1step = 1M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 240.8523 ms 1.1376 ms
STEP2
這一步我們採用分塊tile。分塊的目的是爲了充分利用緩存。如果原來的循環較大，tile分塊改成小塊數據去計算，可以使得每次計算的數據都比較舒適地呆在緩存裏，不用經歷重複的驅逐（在緩存中重複的添加和刪除數據)。分塊後進行reorder操作，交換兩個嵌套循環的順序，目的是最內層的內存訪問友好。我們按照x,y維度劃分成16x8的小分塊去計算:
.gemm.update() .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo);
執行結果如下：
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 2step = 2M N K = 640 640 640 err 0.00 [rep 50] halide | blas 81.8148 ms 1.1281 ms
性能從240ms優化到82ms，提升了近3倍。
STEP3
我們在上一步的基礎上增加向量化vectorize。向量化是把幾個標量計算（scale)轉換爲一個向量計算（vector),充分利用SIMD向量指令。大部分現代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數據流）。在同一個CPU循環中，SIMD可在多個值上同時執行相同的運算/指令。
gemm.update() .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo) .vectorize(xi, 8);
執行結果：
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 3step = 3M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 27.5433 ms 1.1445 ms
性能從82ms優化到27ms，又加速了接近3倍。可以看到，圍繞前面提到的兩條優化宗旨：優化內存訪問和提高並行性，從step1到step3，性能已經提升了近9倍。
STEP4
調度策略在step3的基礎上增加並行化parallel。對一個循環並行化是把循環的每次迭代分給多個線程或者處理器去同時處理，每個線程處理通過代碼段（loop body),但是處理不同的數據。
gemm(x, y) += A(k, y) * B(x, k); gemm.update() .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo) .vectorize(xi, 8) .parallel(yo);
執行結果：
root@bd3faab0f079:/home/chunying/AutoKernel/doc/tutorials# ./06_build.sh 4step = 4M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 7.2605 ms 1.1605 ms
增加並行化後，build.sh默認指定四線程，性能直接翻了近4倍，從27ms到7.3ms.
STEP5
調度策略在上一步的基礎上增加unroll展開。如果循環體內的語句沒有數據相關依賴，循環展開可以增加併發執行的機會，使得更充分利用寄存器，減少循環時每個操作內存加載和保存的次數。
gemm.update() .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo) .vectorize(xi, 8) .parallel(yo) .unroll(xi) .unroll(yi,2);
執行結果：
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 5step = 5M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 4.7617 ms 1.1597 ms
unroll展開後，性能從7.3ms優化到4.8ms.
STEP6
前面的分塊成 16 x 8的小kernel, 這一步先劃分成 16 x 32的分塊，然後把每個分塊再分成 16 x 8的子分塊。我們把最外層的兩層循環合併到一層，並對這一層進行並行化。這一步計算描述多了一個prod函數來定義子分塊的計算，prod函數的計算公式和總的gemm是一樣的，我們通過 compute_at指定在 yi維度之下計算prod，則prod計算的是 16x8的小kernel, 大致邏輯如下：
總的代碼如下：
Func prod; prod(x, y) += A(k, y) * B(x, k); gemm(x, y) = prod(x, y); gemm.tile(x, y, xi, yi, 16, 32) .fuse(x, y, xy).parallel(xy) .split(yi, yi, yii, 4) .vectorize(xi, 8) .unroll(xi) .unroll(yii); prod.compute_at(gemm, yi) .vectorize(x, 8).unroll(y); prod.update() .reorder(x, y, k) .vectorize(x, 8) .unroll(x) .unroll(y) .unroll(k, 2);
執行結果
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 6step = 6M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 3.1824 ms 1.1373 ms
這一步距離STEP1性能已經優化了近80倍了，性能越來越接近OpenBlas了。
STEP 7
這一步添加的操作是對矩陣B進行數據重排，使得在計算小kernel 16x8時，內存讀取更順暢。因爲小kernel的x維度是按照16劃分的，因此重排數據B的x維度也是按照16重排。
總的代碼如下：
Func B_interleave("B"), Bs("Bs"); Bs(x, y, xo) = B(xo * 16 + x, y); B_interleave(x, y) = Bs(x % 16, y, x / 16); Func prod; prod(x, y) += A(k, y) * B_interleave(x, k); gemm(x, y) = prod(x, y); gemm.tile(x, y, xi, yi, 16, 32)
.fuse(x, y, xy).parallel(xy) .split(yi, yi, yii, 4) .vectorize(xi, 8) .unroll(xi) .unroll(yii); prod.compute_at(gemm, yi) .vectorize(x, 8).unroll(y); prod.update() .reorder(x, y, k) .vectorize(x, 8) .unroll(x) .unroll(y) .unroll(k, 2); Bs.compute_root() .split(y, yo, yi, 16) .reorder(x, yi, xo, yo) .unroll(x)
.vectorize(yi).parallel(yo, 4);
執行結果：
root@bd3faab0f079:/AutoKernel/doc/tutorials/data# ./06_build.sh 7step = 7M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 1.1957 ms 1.1425 ms
至此，我們的每一步調優策略始終都圍繞兩條優化宗旨“優化內存訪問”，“提高並行性”展開優化，到最後性能已經與OpenBlAS差不多了，距離STEP1已經加速了200+倍了。
項目地址：
https://github.com/OAID/AutoKernel/
更多精彩推薦