Halide應(yīng)用開發(fā)

1. 基本原理
   1.1、介紹
   隨著人工智能的普及，深度學(xué)習(xí)網(wǎng)絡(luò)的不斷涌現(xiàn)，為了讓各硬件(CPU, GPU, NPU,…)能夠支持深度學(xué)習(xí)應(yīng)用，各硬件芯片需要軟件庫去支持高性能的深度學(xué)習(xí)張量運(yùn)算。目前，這些高性能計(jì)算庫主要由資深HPC工程師(高性能計(jì)算優(yōu)化工程師）進(jìn)行開發(fā)，為了加快開發(fā)進(jìn)程，縮短深度學(xué)習(xí)應(yīng)用落地周期，自動(dòng)化算子優(yōu)化是一個(gè)趨勢(shì)。
   

AutoKernel是由OPEN AI LAB提出的高性能算子自動(dòng)優(yōu)化工具，可以自動(dòng)優(yōu)化調(diào)度策略、生成底層優(yōu)化代碼，大幅減少各硬件芯片算子開發(fā)成本，提升算子優(yōu)化效率，讓工程師更快實(shí)現(xiàn)深度學(xué)習(xí)算法在各硬件芯片上的高性能部署。
1.2、 AutoKernel架構(gòu)

AutoKernel分為3個(gè)模塊：
算子生成器
該模塊使用了開源項(xiàng)目Halide；Halide是業(yè)界廣泛使用的自動(dòng)代碼生成項(xiàng)目，它首次提出將計(jì)算和調(diào)度分離。該模塊的輸入是和硬件無關(guān)的算子計(jì)算描述，輸出是相應(yīng)后端的優(yōu)化匯編代碼/目標(biāo)文件；
自動(dòng)搜索模塊
該模塊可以通過最優(yōu)化算法/搜索算法/機(jī)器學(xué)習(xí)/強(qiáng)化學(xué)習(xí)搜索出相應(yīng)后端的最優(yōu)算子的調(diào)度策略參數(shù)（該模塊仍在開發(fā)中）；
算子部署插件(AutoKernel Plugin)
Tengine是OPEN AILAB開源的深度學(xué)習(xí)推理框架，實(shí)現(xiàn)了AI算法在不同硬件的快速高效部署。該模塊實(shí)現(xiàn)了將自動(dòng)生成的優(yōu)化算子代碼以plugin的形式一鍵集成到Tengine中，實(shí)現(xiàn)自動(dòng)優(yōu)化算子的一鍵部署；
1.3、 AutoKernel特點(diǎn)
? 低門檻: 無需底層優(yōu)化匯編的知識(shí)門檻
? 簡單易用: 提供docker環(huán)境，無需安裝環(huán)境，plugin一鍵集成到推理框架Tengine
? 高效率: 無需手寫優(yōu)化匯編，一鍵生成優(yōu)化代碼，一鍵部署
2 AutoKernel優(yōu)化GEMM實(shí)踐
本部分將帶領(lǐng)大家一步步優(yōu)化矩陣乘法GEMM。無需手工擼代碼，編寫繁雜冗長的底層匯編代碼，只需十幾行簡潔的調(diào)度代碼。

2.1、優(yōu)化的本質(zhì)
在詳細(xì)講解優(yōu)化步驟前，先談?wù)剝?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)：主存容量大，訪問速度慢，寄存器和緩存讀取速度快，但容量有限。在寄存器的層級(jí)上，CPU可以在一個(gè)時(shí)鐘周期內(nèi)訪問它們，如果CPU去訪問外部的DDR的話，延遲是非常大的，大概是200個(gè)時(shí)鐘周期左右。如果CPU去訪問cache的話，一般需要6到12個(gè)cycle就夠了。所以，第1個(gè)優(yōu)化宗旨是：優(yōu)化內(nèi)存訪問，充分利用寄存器和高速緩存去存數(shù)據(jù)。
第2個(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倍的加速。
2.2、運(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.git
cd 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)化的。
執(zhí)行指令：

執(zhí)行demo

./build.sh 1 ==> 默認(rèn)不優(yōu)化
./build.sh 7 ==> 最極致優(yōu)化
下圖展示了在Intel? Core? i9-9900K CPU @ 3.60GHz的電腦上的優(yōu)化效果，無需手工擼代碼，無需編寫繁雜冗長的底層匯編代碼，只需十幾行簡潔的調(diào)度代碼, 就能性能優(yōu)化200+倍~

2.3、詳細(xì)的優(yōu)化步驟：
? STEP-1
第1個(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@xxxxxxxxxx:/AutoKernel/doc/tutorials/data# ./06_build.sh 1
step = 1
M N K = 640 640 640 err 0.00 [rep 50] autokernel | blas 240.8523 ms 1.1376 ms
? STEP-2
這一步采用分塊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@xxxxxxxxxx:/AutoKernel/doc/tutorials/data#./06_build.sh 2
step = 2
M N K = 640 640 640 err 0.00 [rep 50] halide | blas 81.8148 ms 1.1281 ms
性能從240ms優(yōu)化到82ms，提升了近3倍。
? STEP-3
在上一步的基礎(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@xxxxxxxxxx:/AutoKernel/doc/tutorials/data# ./06_build.sh 3
step = 3
M 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倍。
? STEP-4
調(diào)度策略在step3的基礎(chǔ)上增加并行化parallel。對(duì)一個(gè)循環(huán)并行化是把循環(huán)的每次迭代分給多個(gè)線程或者處理器去同時(shí)處理，每個(gè)線程處理通過代碼段（loop body),但是處理不同的數(shù)據(jù)。
增加并行化后，build.sh默認(rèn)指定四線程，性能直接翻了近4倍，從27ms到7.3ms.
? STEP-5
調(diào)度策略在上一步的基礎(chǔ)上增加unroll展開。如果循環(huán)體內(nèi)的語句沒有數(shù)據(jù)相關(guān)依賴，循環(huán)展開可以增加并發(fā)執(zhí)行的機(jī)會(huì)，使得更充分利用寄存器，減少循環(huán)時(shí)每個(gè)操作內(nèi)存加載和保存的次數(shù)。
unroll展開后，性能從7.3ms優(yōu)化到4.8ms.
? STEP-6
前面的分塊成 16 x 8的小kernel, 這一步先劃分成 16 x 32的分塊，然后把每個(gè)分塊再分成 16 x 8的子分塊。把最外層的兩層循環(huán)合并到一層，并對(duì)這一層進(jìn)行并行化。這一步計(jì)算描述多了一個(gè)prod函數(shù)來定義子分塊的計(jì)算，prod函數(shù)的計(jì)算公式和總的gemm是一樣的，通過 compute_at指定在 yi維度之下計(jì)算prod，則prod計(jì)算的是 16x8的小kernel, 大致邏輯如下：

這一步距離STEP1性能已經(jīng)優(yōu)化了近80倍了，性能越來越接近OpenBlas了。
? STEP-7
這一步添加的操作是對(duì)矩陣B進(jìn)行數(shù)據(jù)重排，使得在計(jì)算小kernel 16x8時(shí)，內(nèi)存讀取更順暢。因?yàn)樾ernel的x維度是按照16劃分的，因此重排數(shù)據(jù)B的x維度也是按照16重排。

至此，的每一步調(diào)優(yōu)策略始終都圍繞兩條優(yōu)化宗旨“優(yōu)化內(nèi)存訪問”，“提高并行性”展開優(yōu)化，到最后性能已經(jīng)與OpenBlAS差不多了，距離STEP1已經(jīng)加速了200+倍了。

總結(jié)

以上是生活随笔為你收集整理的Halide应用开发的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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