GEMM與AutoKernel算子優(yōu)化
隨著AI技術(shù)的快速發(fā)展，深度學(xué)習(xí)在各個領(lǐng)域得到了廣泛應(yīng)用。深度學(xué)習(xí)模型能否成功在終端落地應(yīng)用，滿足產(chǎn)品需求，一個關(guān)鍵的指標(biāo)就是神經(jīng)網(wǎng)絡(luò)模型的推理性能。一大波算法工程師為了算法的部署轉(zhuǎn)崗算子優(yōu)化工程師。優(yōu)化代碼并不是一件簡單的事，要求工程師既要精通計算機體系架構(gòu)，又要熟悉算法的計算流程，稍微有經(jīng)驗的深度學(xué)習(xí)推理優(yōu)化工程師都成了各家公司爭搶的“香餑餑”。需求多，算子優(yōu)化自動化成為了未來的一大趨勢。

為了方便更多的工程師進行推理優(yōu)化，一個致力于降低優(yōu)化門檻，提升優(yōu)化開發(fā)效率的算子自動優(yōu)化工具AutoKernel宣布正式開源！

AutoKernel特色：
低門檻: 無需底層優(yōu)化匯編的知識門檻簡單易用: 提供docker環(huán)境，無需安裝環(huán)境，plugin一鍵集成到推理框架Tengine高效率: 無需手寫優(yōu)化匯編，一鍵生成優(yōu)化代碼，一鍵部署AutoKernel使用業(yè)界廣泛使用的自動代碼生成項目Halide，通過輸入計算描述和調(diào)度策略，自動生成底層代碼。AutoKernel支持以plugin的形式，將生成的自動優(yōu)化算子一鍵部署到推理框架Tengine中。
本文將帶領(lǐng)大家一步步優(yōu)化矩陣乘法GEMM。無需手工擼代碼，編寫繁雜冗長的底層匯編代碼，只需十幾行簡潔的調(diào)度代碼。

優(yōu)化的本質(zhì)。優(yōu)化的時候，計算機底層做了什么？優(yōu)化的”瓶頸“是什么？為什么通過一波”優(yōu)化操作“，性能就能提升呢？AutoKernel使用的Halide是如何實現(xiàn)自動優(yōu)化的呢？
需要了解一下硬件的基礎(chǔ)的體系結(jié)構(gòu)，了解硬件如何工作，才能在軟件上實現(xiàn)算法的時候，盡可能去考慮利用硬件的一些特性，來做到高效的、極致的優(yōu)化。

上圖是典型的存儲理器層次結(jié)構(gòu)：主存容量大，訪問速度慢，寄存器和緩存讀取速度快，但容量有限。在寄存器的層級上，CPU可以在一個時鐘周期內(nèi)訪問它們，如果CPU去訪問外部的DDR的話，延遲是非常大的，大概是200個時鐘周期左右。如果CPU去訪問cache的話，一般需要6到12個cycle就夠了。所以，一個很重要的一個優(yōu)化宗旨是：優(yōu)化內(nèi)存訪問，充分利用寄存器和高速緩存去存數(shù)據(jù)。
第二個優(yōu)化宗旨則是提高并行性：充分利用SIMD進行指令向量化和多核心并行。大部分現(xiàn)代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數(shù)據(jù)流）。在同一個CPU循環(huán)中，SIMD可在多個值上同時執(zhí)行相同的運算/指令。在4個數(shù)據(jù)點上進行向量化，一次計算四個數(shù)據(jù)，理論上就可以實現(xiàn)4倍的加速。

運行環(huán)境搭建
AutoKernel提供了docker鏡像，docker里已經(jīng)配置好運行環(huán)境，進入docker即可直接運行demo代碼：

拉取鏡像docker pull openailab/autokernel# 啟動容器，進入開發(fā)環(huán)境docker run -it openailab/autokernel /bin/bash# 獲取代碼git clone https://github.com/OAID/AutoKernel.gitcd AutoKernel/doc/tutorials/data/目錄下的build.sh是demo的執(zhí)行腳本，運行需要指定優(yōu)化步驟step，可選的step是從1 到7，其中step= 1 是默認(rèn)不優(yōu)化的，step=7是最極致優(yōu)化的。

優(yōu)化效果

執(zhí)行demo./build.sh 1./build.sh 7

下圖展示了在Intel? Core? i9-9900K CPU @ 3.60GHz的電腦上的優(yōu)化效果，無需手工擼代碼，無需編寫繁雜冗長的底層匯編代碼，只需十幾行簡潔的調(diào)度代碼, 就能性能優(yōu)化200+倍~

優(yōu)化步驟
以下是更為詳細(xì)的優(yōu)化步驟：
STEP1
第一個步驟是不帶任何優(yōu)化的。用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的矩陣乘法。運行腳本第一個參數(shù)指定step=1。耗時結(jié)果如下：
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。分塊的目的是為了充分利用緩存。如果原來的循環(huán)較大，tile分塊改成小塊數(shù)據(jù)去計算，可以使得每次計算的數(shù)據(jù)都比較舒適地呆在緩存里，不用經(jīng)歷重復(fù)的驅(qū)逐（在緩存中重復(fù)的添加和刪除數(shù)據(jù))。分塊后進行reorder操作，交換兩個嵌套循環(huán)的順序，目的是最內(nèi)層的內(nèi)存訪問友好。按照x,y維度劃分成16x8的小分塊去計算:
.gemm.update .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo);
執(zhí)行結(jié)果如下：
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優(yōu)化到82ms，提升了近3倍。
STEP3
在上一步的基礎(chǔ)上增加向量化vectorize。向量化是把幾個標(biāo)量計算（scale)轉(zhuǎn)換為一個向量計算（vector),充分利用SIMD向量指令。大部分現(xiàn)代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數(shù)據(jù)流）。在同一個CPU循環(huán)中，SIMD可在多個值上同時執(zhí)行相同的運算/指令。
gemm.update .tile(x, y, xo, yo, xi, yi, 16, 8) .reorder(xi, yi, k, xo, yo) .vectorize(xi, 8);
執(zhí)行結(jié)果：
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優(yōu)化到27ms，又加速了接近3倍?？梢钥吹?#xff0c;圍繞前面提到的兩條優(yōu)化宗旨：優(yōu)化內(nèi)存訪問和提高并行性，從step1到step3，性能已經(jīng)提升了近9倍。
STEP4
調(diào)度策略在step3的基礎(chǔ)上增加并行化parallel。對一個循環(huán)并行化是把循環(huán)的每次迭代分給多個線程或者處理器去同時處理，每個線程處理通過代碼段（loop body),但是處理不同的數(shù)據(jù)。
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);
執(zhí)行結(jié)果：
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默認(rèn)指定四線程，性能直接翻了近4倍，從27ms到7.3ms.
STEP5
調(diào)度策略在上一步的基礎(chǔ)上增加unroll展開。如果循環(huán)體內(nèi)的語句沒有數(shù)據(jù)相關(guān)依賴，循環(huán)展開可以增加并發(fā)執(zhí)行的機會，使得更充分利用寄存器，減少循環(huán)時每個操作內(nèi)存加載和保存的次數(shù)。
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);
執(zhí)行結(jié)果：
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優(yōu)化到4.8ms.
STEP6
前面的分塊成 16 x 8的小kernel, 這一步先劃分成 16 x 32的分塊，然后把每個分塊再分成 16 x 8的子分塊。把最外層的兩層循環(huán)合并到一層，并對這一層進行并行化。這一步計算描述多了一個prod函數(shù)來定義子分塊的計算，prod函數(shù)的計算公式和總的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);
執(zhí)行結(jié)果
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性能已經(jīng)優(yōu)化了近80倍了，性能越來越接近OpenBlas了。
STEP 7
這一步添加的操作是對矩陣B進行數(shù)據(jù)重排，使得在計算小kernel 16x8時，內(nèi)存讀取更順暢。因為小kernel的x維度是按照16劃分的，因此重排數(shù)據(jù)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);
執(zhí)行結(jié)果：
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
至此，的每一步調(diào)優(yōu)策略始終都圍繞兩條優(yōu)化宗旨“優(yōu)化內(nèi)存訪問”，“提高并行性”展開優(yōu)化，到最后性能已經(jīng)與OpenBlAS差不多了，距離STEP1已經(jīng)加速了200+倍了。

總結(jié)

以上是生活随笔為你收集整理的GEMM与AutoKernel算子优化的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。