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

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

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

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

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

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

拉取鏡像docker pull openailab/autokernel# 啟動(dòng)容器，進(jìn)入開發(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ùn)行需要指定優(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
第一個(gè)步驟是不帶任何優(yōu)化的。用Halide語言直接描述GEMM的計(jì)算過程。
Var x,y; RDom k(0, K); Func gemm(“gemm”); gemm(x, y) += A(k, y) * B(x, k);
計(jì)算M=N=K=640的矩陣乘法。運(yùn)行腳本第一個(gè)參數(shù)指定step=1。耗時(shí)結(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ù)去計(jì)算，可以使得每次計(jì)算的數(shù)據(jù)都比較舒適地呆在緩存里，不用經(jīng)歷重復(fù)的驅(qū)逐（在緩存中重復(fù)的添加和刪除數(shù)據(jù))。分塊后進(jìn)行reorder操作，交換兩個(gè)嵌套循環(huán)的順序，目的是最內(nèi)層的內(nèi)存訪問友好。按照x,y維度劃分成16x8的小分塊去計(jì)算:
.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。向量化是把幾個(gè)標(biāo)量計(jì)算（scale)轉(zhuǎn)換為一個(gè)向量計(jì)算（vector),充分利用SIMD向量指令。大部分現(xiàn)代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數(shù)據(jù)流）。在同一個(gè)CPU循環(huán)中，SIMD可在多個(gè)值上同時(shí)執(zhí)行相同的運(yùn)算/指令。
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倍。可以看到，圍繞前面提到的兩條優(yōu)化宗旨：優(yōu)化內(nèi)存訪問和提高并行性，從step1到step3，性能已經(jīng)提升了近9倍。
STEP4
調(diào)度策略在step3的基礎(chǔ)上增加并行化parallel。對一個(gè)循環(huán)并行化是把循環(huán)的每次迭代分給多個(gè)線程或者處理器去同時(shí)處理，每個(gè)線程處理通過代碼段（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í)行的機(jī)會，使得更充分利用寄存器，減少循環(huán)時(shí)每個(gè)操作內(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的分塊，然后把每個(gè)分塊再分成 16 x 8的子分塊。把最外層的兩層循環(huán)合并到一層，并對這一層進(jìn)行并行化。這一步計(jì)算描述多了一個(gè)prod函數(shù)來定義子分塊的計(jì)算，prod函數(shù)的計(jì)算公式和總的gemm是一樣的，通過 compute_at指定在 yi維度之下計(jì)算prod，則prod計(jì)算的是 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進(jìn)行數(shù)據(jù)重排，使得在計(jì)算小kernel 16x8時(shí)，內(nèi)存讀取更順暢。因?yàn)樾ernel的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)容還不錯(cuò)，歡迎將生活随笔推薦給好友。