CUDA卷積算子手寫詳細(xì)實現(xiàn)
本文經(jīng)自動駕駛之心公眾號授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
前言
CUDA介紹(from chatGPT)
現(xiàn)在深度學(xué)習(xí)大行其道,作為深度學(xué)習(xí)的基礎(chǔ)軟件設(shè)施,學(xué)習(xí)cuda也是很有意義的。本篇文章主要介紹如何利用CUDA實現(xiàn)一個2D卷積算子,實現(xiàn)過程較為簡單,最終的實現(xiàn)效果可以在較小的尺寸下取得比cudnn快較大的性能。實測在以下參數(shù)配置下可以達(dá)到平均1.2倍cudnn的性能(娛樂結(jié)果,還與cudnn配置有關(guān),更小更快)。
TIPS: 跳過cudnn初始化的時間,99輪平均時間
const int inC = 6;
const int inH = 768;
const int inW = 512;
const int kernelH = 6;
const int kernelW = 6;
const int outC = 6;
const int outH = inH - kernelH + 1;
const int outW = inW - kernelW + 1;
1 卷積操作通俗介紹
1.1 數(shù)據(jù)布局(data layout)
卷積操作主要針對圖像進(jìn)行運算,我們常見的RGB即為三通道的二維圖像,那么就可以通過一個一維數(shù)組存儲所有的數(shù)據(jù),再按照不同的布局去索引對應(yīng)的數(shù)據(jù),現(xiàn)在主要使用nchw和nhwc兩種數(shù)據(jù)布局,其中
n - batch size 也可以理解為"圖像"數(shù)量
c - channel num 即我們說的通道數(shù)量
h - height 圖像高度,每個通道的高度寬度是一致的
w - width 圖像寬度
那么顯然nchw就是逐個通道的讀取圖像,nhwc即對所有通道的同樣位置讀取數(shù)據(jù)后,再切換到下一個為止
一個是優(yōu)先通道讀取,一個是優(yōu)先位置讀取
還有一種CHWN布局,感覺比較奇怪,并未過多了解
詳細(xì)的可以參考英偉達(dá)官方文檔Developer Guide : NVIDIA Deep Learning cuDNN Documentation (https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html)
nchw layout
nhwc layout
本文是按照nchw數(shù)據(jù)格式來進(jìn)行算子的實現(xiàn)的。
1.2 直接卷積
相信大家都或多或少聽過卷積,可以通過gpt的回答來直觀地認(rèn)識卷積操作
最基本的直接卷積操作是十分簡單的,你可以想象一個滑動的矩陣窗口在原矩陣上移動,對應(yīng)位置進(jìn)行點積,得到結(jié)果后求和放到目標(biāo)矩陣上,可以用以下圖像直觀地理解這一過程,向老師稱為對對碰:)
圖源:國科大模式識別課程
你會注意到上述過程中怎么沒有什么channel的參與,只有一個矩陣
多輸入通道的情況下,就是對每個通道的相同位置分別與卷積核進(jìn)行對對碰,結(jié)果累加作為輸出矩陣值;
多輸入多輸出通道,即對每個輸出通道都進(jìn)行上述操作
對于通道的理解建議參考[@雙手插袋]的文章CNN卷積核與通道講解 (https://zhuanlan.zhihu.com/p/251068800)
那么我們需要知道的是直接卷積操作其實就是原矩陣與卷積核間的對對碰,產(chǎn)生所謂的特征圖feature map,十分的簡單,這也方便我們對其進(jìn)行并行任務(wù)劃分
注意到上述文章中并沒有提到padding和stride,本篇文章并沒有針對padding和stride的實現(xiàn)
padding
padding是作為對圖像的填充,可以發(fā)現(xiàn)上面的特征圖尺寸縮小了一圈,是因為直接卷積勢必會造成這一結(jié)果
通過padding可以加強(qiáng)圖像邊緣特征,避免邊緣特征被忽略
stride
stride可以簡單的理解為跨步,即上面的小窗口在矩陣上滑動的步長,默認(rèn)為1
即上述圖像中下一次卷積的中心應(yīng)該是4為中心的3*3子矩陣
如果你設(shè)置為2,那么下一次是3為中心的3*3子矩陣了
1.3 其他卷積計算方法
除去直接卷積,也有一些其他方法進(jìn)行卷積,感興趣的讀者可以自行了解,僅舉以下幾例參考
Img2col
即把圖像展開為一個行向量組,卷積核/濾波器(kernel/filter)展開為一列或多列向量,轉(zhuǎn)化為矩陣乘去計算卷積結(jié)果
FFT method
利用傅里葉變換的頻域變換去做卷積,這樣做的優(yōu)勢是計算量會小很多
Winograd Algorithm
也是一種將圖像變換到另外一個空間再去做運算再做變換得到結(jié)果,會減少很多乘法運算
2 整體實現(xiàn)思路
2.1 block與thread劃分
首先我們需要考慮如何對代表圖像的多通道矩陣來進(jìn)行block與thread的劃分,這一部分是有說法的
不同的切分方式會讓block在SM上的流轉(zhuǎn)效率有很大的差別
本文僅提供一個十分草率的切分,我們都清楚目前在英偉達(dá)的GPU上,任務(wù)的調(diào)度最小單元是warp
一個warp以32個線程為一組,故通過8*4的block來進(jìn)行矩陣的切分,每個block里共32個位置
這樣可以保證每個block上到SM時不用去與其他的block拼接線程,產(chǎn)生額外開銷
注意我這里用的是位置,并不是元素,32個線程,每個線程去負(fù)責(zé)一個位置的計算
以16*20的矩陣為例,對其進(jìn)行劃分的結(jié)果如下圖所示,(x,y)是笛卡爾坐標(biāo)系,與行主序不同
2.2 數(shù)據(jù)轉(zhuǎn)移
- 關(guān)于位置和規(guī)模(size)
那么為什么說一個block有32個位置,而不是32個元素呢,首先注意到卷積操作雖然遍歷到了原矩陣的所有元素
但是你按中心點的位序去數(shù)的話(以卷積核3*3為例),結(jié)果應(yīng)該是這個樣子
注意這里僅示意卷積中心點范圍,請與后文作區(qū)分
按3*3矩陣的中心來看,中心正好是去掉外面一圈的位置,按照左上角元素來看,恰好應(yīng)該是(左上角,右下角)
這樣一個區(qū)間,參數(shù)解釋如下
row_num 原矩陣中一行元素的數(shù)目
inH inW 原矩陣的H W
kernelH kernelW 卷積核的H W
outH outW 輸出矩陣的H W
當(dāng)然你也可以用中心點而不是左上角的元素作為窗口的標(biāo)識來設(shè)計算法
恰巧你上面算出來的這個范圍也正是你得到的feature map的下標(biāo)范圍
我們也可以得到輸出矩陣的規(guī)模為
請注意大小和位置下標(biāo)的區(qū)別,一個從1開始數(shù)一個從0開始數(shù)
- 一個block的數(shù)據(jù)轉(zhuǎn)移
確定了整體的尺寸,那么我們來看一個block需要的數(shù)據(jù)尺寸是多少呢
顯然你可以發(fā)現(xiàn),對于輸出矩陣進(jìn)行block劃分是更合理的,這樣可以保證一個block
32個位置恰好對應(yīng)輸出矩陣的32個位置,而不用過多的去考慮輸出矩陣的排布
那么對于上述提到的劃分,可以通過下圖來直觀感受block劃分在原矩陣的效果
22*18的in產(chǎn)生20*16的out
那么一個block用到的元素范圍應(yīng)該是哪些呢,我們要做的是卷積操作,每個中心點應(yīng)該有對應(yīng)卷積核大小的矩陣參與運算,那么以(0,0)和(4,1)的block為例,給出他們的涉及原矩陣范圍如下圖所示
藍(lán)色為一個block需要用到的原矩陣元素
那么我們可以確定一個block,8×4的情況下需要讀取10×6的原矩陣的元素,也是+kernelH-1來確定的
那么對應(yīng)輸出矩陣就是一個蘿卜一個坑了,不需要額外考慮
這樣就確定了一個block需要從GMEM到SMEM的元素范圍
至于怎么轉(zhuǎn)移,我們在代碼實現(xiàn)中講述,當(dāng)然你可以單獨指定某幾個進(jìn)程去完成所有的轉(zhuǎn)移任務(wù)
2.3 計算邏輯
- 不考慮channel
不考慮channel的情況下,即單輸入通道單輸出通道單卷積核這樣最簡單的情況
我們只需要做三件事
① 將block對應(yīng)的數(shù)據(jù)轉(zhuǎn)移到SMEM中
② 利用線程的tid去計算對應(yīng)輸出矩陣位置的結(jié)果
③ 將結(jié)果寫回輸出矩陣
- 只考慮inC
這種情況下我們要做的額外的事兒就多一點
加一層循環(huán),讓每個線程計算多個in channel的數(shù)據(jù),并累加起來作為結(jié)果
需要用到一個寄存器來存儲這個中間結(jié)果
- 考慮inC與outC
其實要做的事情也就比上面多一點,就是開大點空間
讓線程去存儲多個outC的中間結(jié)果,分別累加
最后寫回的時候也分別寫回即可
3 詳細(xì)實現(xiàn)過程
3.1 整體實現(xiàn)思路
主要從自己的角度出發(fā)去還原怎樣一步步構(gòu)造出這樣一個初級的算法
首先實現(xiàn)一個最簡單的版本,CPU串行版本,并保證CPU串行版本可以獲取正確的結(jié)果
此后再在其基礎(chǔ)上進(jìn)行并行化的改造,而直接卷積運算的過程其實相對是比較簡單的
我們在不考慮padding與stride的情況下,是可以不借助任何參考資料來直接完成第一版代碼的
3.1.1 CPU串行版本的卷積算子
#define element_type float
#define OFFSET(row, col, ld) ((row) * (ld) + (col))
/*
@brief: 串行卷積實現(xiàn) CPU代碼 NCHW
@param in inC inH inW: 輸入矩陣(數(shù)組) channel height width
@param out outC outH outW: 輸出矩陣 channel height width
@param kernel kernelH kernelW: 卷積核 height width
*/
void serial_convolution(element_type *in, element_type *out, element_type *kernel, int batch_size,
int inC, int inH, int inW,
int outC, int outH, int outW,
int kernelH, int kernelW)
{
float val;
int out_pos, in_pos, kernel_pos;
for (int oc = 0; oc < outC; oc++) // 每個輸出通道
{
// 對一個位置的操作 用當(dāng)前輸入channel卷積去對相應(yīng)的輸出channel
// 保證每個outChannel都是多inChannel累積的結(jié)果
for (int i = 0; i < outH; i++)
{
for (int j = 0; j < outW; j++)
{
val = 0; // 避免累積和需要多次讀取寫入
out_pos = oc * outH * outW + OFFSET(i, j, outW);
for (int ic = 0; ic < inC; ic++) // 對每個輸入通道
{
for (int ii = 0; ii < kernelH; ii++)
{
for (int jj = 0; jj < kernelW; jj++)
{
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);
val += in[in_pos] * kernel[kernel_pos];
}
}
}
out[out_pos] = val; // 與cudnn計算結(jié)果為相反數(shù)
}
}
}
}
這是我最終完成的CPU串行版本代碼,可以發(fā)現(xiàn)套了足足有5層循環(huán)
在我們傳統(tǒng)觀念中,這可是 O(n5)O(n^5)O(n^5) 的最笨算法了
不過沒有關(guān)系,我們關(guān)注的并不是他的性能,cuda上也不會去跑這一版代碼
我們需要關(guān)注的是怎么樣能得到正確的結(jié)果,且如何設(shè)計循環(huán)的嵌套關(guān)系來使用盡量少的訪存次數(shù)
使用盡量多的本地中間結(jié)果,這樣可以盡可能地減少我們的算法在訪存方面的消耗
要明白GPU上的線程如果去讀GMEM上的數(shù)據(jù)需要幾百個時鐘周期,讀SMEM需要幾十個時鐘周期
讀取SM上的寄存器需要的時鐘周期會更少!
因此我們需要竭力優(yōu)化的一部分是如何減少訪存,多用本地存儲來代替
另一方面這也是因為計算本身是十分簡單的點積,不太可能去做出更大的優(yōu)化
3.1.2 循環(huán)順序設(shè)計
逐層去觀察循環(huán)的嵌套順序,發(fā)現(xiàn)是
outC-->H-->W--->inC-->kernelH-->kernelW
這樣的計算順序不一定是最優(yōu)化的,筆者也沒有進(jìn)行詳細(xì)的計算論證,但是這樣的計算順序是出于以下角度考慮
① 多通道卷積結(jié)果的維度/通道數(shù)/feature map數(shù)就是我們的outC,是我們最終要寫回的out矩陣的維度,將其放在最外層循環(huán),作用是:
- 一次循環(huán)內(nèi)完成這個out channel中的所有計算,再接著進(jìn)行下一個outC的計算
- 由于out數(shù)據(jù)是在一維數(shù)組中存儲,且為nchw格式,那么不同outC中的數(shù)據(jù)跨度其實是很大的,連續(xù)的完成一個outC的內(nèi)容可以更好的利用局部性原理
- 個人理解逐個outC的計算是很是一種比較直觀和自然(方便想象與理解)的角度
- 串行過程中我們可以使用盡量少的中間變量去維護(hù)中間結(jié)果,如果你先遍歷inC后遍歷outC的話,其實你是需要維護(hù)outC個中間變量的
這樣的順序也是在后面做并行化改造過程中逐漸發(fā)現(xiàn)的一個較為合理的順序,我們可以在后文中更加直觀的感受到這樣設(shè)計的優(yōu)勢
② 出于nchw布局的涉及,H W的順序是基本固定的,當(dāng)然你也可以先W后H,不過一般是行主序存儲.. 還是先H比較快一些
③ inC為何出現(xiàn)在H W之后?請回顧多通道卷積的過程,一個feature map的值是由多個inC與kernel分別點擊累加形成的,如果你將inC放置在H W之前的話,在下方的代碼中,你是不是就需要設(shè)置height×width個中間變量來存儲這里的val值呢?
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);
val += in[in_pos] * kernel[kernel_pos];
將inC放置在H W之后,是相當(dāng)于在一個outC上進(jìn)行計算,對不同inC同樣的位置分別計算得到了val的準(zhǔn)確值,最終寫回,這樣在串行的版本中,我們只需要一個float即可存儲好中間結(jié)果來避免空間的浪費!
TIPS:注意上方對于下標(biāo)的計算,我們以兩個位序舉例說明
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
nchw的數(shù)據(jù)布局格式下,這里是默認(rèn)n為1的,注意本文所有的實現(xiàn)都是建立在n假設(shè)為1的情況,其實n為更大值也不是很有意義,這樣的布局下,下一張圖像在計算意義上是沒有任何差別的,無非是你將數(shù)據(jù)的起始地址跳過一大部分,切到下一張圖像
說回這個式子,其中ic為in channel,inH inW分別是輸入矩陣的高度與寬度,后面宏定義的OFFSET其實就是簡略寫法,你也可以寫成(i+ii)*inW + j + jj
in_pos的含義是在當(dāng)前循環(huán)變量下輸入矩陣的位置
同理,out_pos的計算是一樣的
out_pos = oc * outH * outW + OFFSET(i, j, outW);
ii和jj是相對于卷積核的相對位置循環(huán)變量,輸出位置是用不到他們的
進(jìn)行并行化改造
其實當(dāng)你把串行版本設(shè)計明白后,你對于并行化改造的想法也差不多有個七七八八了
主要是出于以下三個角度去設(shè)計并優(yōu)化的
① 盡量減少訪存次數(shù)(當(dāng)然不是不訪問),尤其是減少訪問GMEM的次數(shù),善用SMEM與register
(對于GMEM SMEM和register等訪存層次相關(guān)知識不熟的讀者可以去了解一下CUDA的存儲層次)
② 此外要劃分明確各個線程要負(fù)責(zé)的任務(wù)區(qū)域和他的行為應(yīng)達(dá)到的效果,做好下標(biāo)計算
③ 計算行為是很快的,我們要盡可能去掩蓋訪存延遲,讓線程去火力全開計算(預(yù)取prefetch)
下面的章節(jié)都是在并行化改造過程中的一些細(xì)節(jié),代碼其實是一版版寫出來的,這里是對最終版本進(jìn)行說明
(所謂的一版版就是劃分出不同塊,分別測試是否與預(yù)期一致,再去完成下面的塊)
3.2 線程任務(wù)均分
這部分其實是源于 @有了琦琦的棍子 在GMEM講解中的數(shù)據(jù)轉(zhuǎn)移部分,基本算是照抄了
十分感謝前輩,不過還不知道這種方法的確切名字,目前暫時稱為均分,其實思想是很樸素的
我們的block設(shè)計的是8*4的大小,對應(yīng)32個線程,但是涉及到in矩陣的數(shù)據(jù)可不只是32個元素,那么
我們需要盡可能地平均分配任務(wù)給線程,保證每個線程承擔(dān)差不多的任務(wù)量來達(dá)到更好的平均性能
差不多是因為,不太可能都是整除的情況
這部分主要通過圖示講解,自己設(shè)計的過程中大多是通過紙筆演算確定下標(biāo)的
首先確定一些變量,注意CUDA的笛卡爾坐標(biāo)系和筆者的行號row和列號col的區(qū)別
int block_row = blockIdx.y;
int block_col = blockIdx.x;
int thread_row = threadIdx.y, thread_col = threadIdx.x;
int tid = thread_row * threadW + thread_col;
由于要重復(fù)使用inC內(nèi)的數(shù)據(jù),我們肯定是要開一個SMEM去存儲這部分?jǐn)?shù)據(jù)的,那么就有一個GMEM->SMEM的數(shù)據(jù)轉(zhuǎn)移過程,以8×4的block和3×3的kernel為例,我們可以得到如下的景象
其中橙色部分是我們的block,一個tid(thread id)是一個線程,也是block中的一個位置,也是outC中的一個位置
那么白色部分就是我們在block范圍之外但會用到的數(shù)據(jù),這部分?jǐn)?shù)據(jù)可以看到像兩條網(wǎng)格
那么我們怎么把這些數(shù)據(jù)從GMEM轉(zhuǎn)移到SMEM呢,首先我們考慮(以下部分為自己笨拙的思考過程)
方案① 邊緣線程負(fù)責(zé)白色區(qū)域
橙色為僅負(fù)責(zé)自己的位置,紫色負(fù)責(zé)3個位置,紅色負(fù)責(zé)9個
看起來是不是好像也還行,只要我們通過thread_row和thread_col判斷一下當(dāng)前進(jìn)程是否在邊緣
對這些進(jìn)程進(jìn)行單獨的編碼就可以了,不過在寫代碼前可以先算一筆賬
這個網(wǎng)格共有10×6=60個元素,我們有32個線程,那么最好的情況下,是每個線程負(fù)責(zé)
60/32=1.875個元素,也就是花費1.875個單位時間(這里的單位時間是抽象概念,假定為每個線程處理每個元素的時間)
那么可以看一下這種劃分方式下,每個線程平均負(fù)責(zé)的元素為
后面的項是權(quán)重,前面的項如 說明這個線程處理9個線程,那么花費的時間應(yīng)當(dāng)是9倍,所以性能應(yīng)當(dāng)是九分之一(相當(dāng)于只處理一個元素的線程),且線程是warp調(diào)度的,32個線程里面有這么一個拖后腿分子,想必并行情況下整體花費時間是取決于這個31號線程的
這個方案的效率是理想情況的一半都不到,說明這種方案是不太可行的,寫出來效果也不一定好呢,換!
方案② 平均劃分
其實筆者也想過一些其他奇怪的方法,但是感覺平均思想似乎是最佳的,那么何不一步到胃呢?
我們先來定義一些變量,后面再來逐步解釋
// 分塊邊界 boundary是限制正常范圍 edge是需要補(bǔ)的范圍
int row_boundary = outH / BLOCK_HEIGHT - 1,
col_boundary = outW / BLOCK_WIDTH - 1;
int row_edge = outH % BLOCK_HEIGHT, col_edge = outW % BLOCK_WIDTH;
···
int single_trans_ele_num = 4; // 線程一次轉(zhuǎn)移的數(shù)據(jù)數(shù)
int cur_in_block_height = BLOCK_HEIGHT + KERNEL_HEIGHT - 1, // 讀入in的block height
cur_in_block_width = BLOCK_WIDTH + KERNEL_WIDTH - 1, // 讀入in的block width
in_tile_thread_per_row, // 以tile為單位轉(zhuǎn)移數(shù)據(jù),一行需要的thread數(shù)
in_tile_row_start, // tile的行起始位置
in_tile_col, // tile的列
in_tile_row_stride; // tile行跨度
// 修正邊緣block尺寸
if (block_row == row_boundary)
{
cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{
cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}
in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;
3.2.1 “block”設(shè)計與修正
不要急著頭大,我們逐個說明,首先看頂頭部分的變量,是關(guān)于限制范圍的
因為我們要首先確定一個block內(nèi)的線程要負(fù)責(zé)多少元素呢,因此需要界定這樣的范圍
我們前面只提到了block涉及到的in范圍是擴(kuò)大了一圈的,其實你的in矩陣相對于out矩陣也是多了一圈的
當(dāng)多的這么一圈不能構(gòu)成新的block時,那么注定我們的block網(wǎng)格是不能覆蓋到out矩陣的!
我們還是上圖比較直觀
咱們的block網(wǎng)格只有16×20這么大,out矩陣有18×22這么大,明顯可以看到藍(lán)色的兩條
是不足以構(gòu)成新的block的,那么還有紅色的部分,就是in矩陣的大小了,可以看到有20×24這么大
而我們的block是建立在out矩陣上的,所以我們起碼也要覆蓋到藍(lán)色矩陣的所有范圍吧
那么在不修改block尺寸的情況下,最簡單的方法就是人為地去修正這些特定block的大小啦
修正后的block應(yīng)該是這個樣子的
修正后的block把out全覆蓋了~
怎么修正呢?無非就是利用block位序去判斷并修改尺寸啦,即這兩行代碼
// 修正邊緣block尺寸
if (block_row == row_boundary)
{
cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{
cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}
結(jié)合圖片,是不是這些變量的概念就清晰了起來
注意我們所有變量都是有一個in的標(biāo)識,這是標(biāo)注in矩陣的范圍
out矩陣的劃分自然是有out的標(biāo)識,且步驟都是一樣的,只不過需要補(bǔ)的范圍不太一樣罷了
3.2.2 線程行為指定
還有一段代碼我們沒有解釋,是這一段(thread_num_per_block本文默認(rèn)為32,沒有修改)
in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;
這段我覺得是最抽象的部分也恰恰是最為精華的設(shè)計,首先要明確,是通過行里面的小片/tile作為線程處理的最小單元來進(jìn)行設(shè)計的
其實變量名已經(jīng)做了一部分的解釋,可以大概解釋為如下的含義
in_tile_thread_per_row 一行里面會有多少個tile
in_tile_row_start 當(dāng)前線程負(fù)責(zé)的tile的起始行號
in_tile_col 當(dāng)前線程負(fù)責(zé)的列號
in_tile_row_stride 如果還有元素要處理,那么需要跳過的行數(shù)/stride
好像不是那么的直觀,我們再上一張圖
左面是我們的block與in矩陣的關(guān)系,我們要把他都轉(zhuǎn)移過來,且利用了fetch_float4的向量指令(也是single_trans_ele_num設(shè)置為4的原因)
以7號線程為例,當(dāng)前的in_block為10×6大小,那么上面四個變量的值分別為1,7,0,32
這個例子比較簡單,可以發(fā)現(xiàn)一行其實是有一個半的tile的,那么需要一點點小小的修正來讓每個線程
讀取4+2個元素,這點小小的修正我們可以看代碼
那么再來一個復(fù)雜的例子,假設(shè)我們在考慮out矩陣的事情,那么一個線程負(fù)責(zé)一個元素的話
請問這種方式對嘛?
是不是直觀上你感覺應(yīng)該是這樣的,他可以絲滑的銜接好每個元素,完成我們的分配~
那么給出我們利用這個均分思想讓每個線程負(fù)責(zé)任務(wù)的代碼如下,大家再想一想分配后的圖像
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
/*do something*/
}
淺淺一個for循環(huán),只不過所有條件都是我們仔細(xì)設(shè)計的,循環(huán)內(nèi)部就是每個線程根據(jù)這些位序
去對應(yīng)的顯存位置上對數(shù)據(jù)一通操作罷了
那么注意部分,線程在跨過一個stride時,這個單位是不是row?那么意味著0號線程在下次任務(wù)會踩到30號的位置!如下圖所示
實際上的線程分配
這樣才是正確的線程操作順序,當(dāng)然由于我們是通過CUDA并行計算的,實際上上半部分是并行的,下半部分是在0-29號線程完成了上面的任務(wù)后才進(jìn)行計算的(注意他們是32個一組/warp調(diào)度上來執(zhí)行的)
這樣其實有個小隱患,30號和31號以及0,1號會對這兩個位置上重復(fù)進(jìn)行操作,如果他們的行為不一致的話
會導(dǎo)致我們的結(jié)果出錯,本例中他們的行為是一致的,故無所謂先后
通過這樣的機(jī)制,我們可以指定每個線程負(fù)責(zé)的元素位置以及個數(shù)(tile大小),靈活地應(yīng)用于不同的任務(wù)!
3.3 預(yù)取機(jī)制
這部分就是很基本的數(shù)據(jù)預(yù)取,計算的效率遠(yuǎn)遠(yuǎn)大于訪存,計算時讀取數(shù)據(jù)進(jìn)來,完成基本的運算
(復(fù)雜運算也不是一行代碼可以解決的)
再把結(jié)果存到對應(yīng)位置,我們發(fā)現(xiàn)是不是即使是計算你也需要訪存,節(jié)省訪存開銷是十分重要的
整體的數(shù)據(jù)傳輸邏輯是GMEM->SMEM->register->GMEM->MEM
并沒有使用到Constant Memory和Texture Memory,那么結(jié)合數(shù)據(jù)預(yù)取的機(jī)制下
整體的框架如下方偽代碼所示
初始化我們所需要的所有變量并修正block規(guī)模;
分配好shared memory用于加速訪存;
// 預(yù)讀取第一個channel的數(shù)據(jù)
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
把in中的數(shù)據(jù)從GMEM轉(zhuǎn)到SMEM;
}
// 預(yù)讀取第一個kernel的數(shù)據(jù) 可以使用很簡單的讀取策略 因為數(shù)據(jù)很少
if (thread_row >= 0 && thread_row < KERNEL_HEIGHT && thread_col == 0)
{
把kernel的數(shù)據(jù)從GMEM轉(zhuǎn)到SMEM;
}
__syncthreads();
// 這里oc在外ic在內(nèi)的設(shè)計是為了方便寫回
for (int oc = 0; oc < outC; oc++)
{
for (int ic = 0; ic < inC; ic++)
{
// i,j 是相當(dāng)于當(dāng)前block起始位置而言
// 用ic的每個block去對oc的kernel進(jìn)行計算
for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height;
i += out_tile_row_stride)
{
計算當(dāng)前ic與oc的結(jié)果,存到register;
}
// 讀取下一個in channel數(shù)據(jù) 3,932,160
if (ic + 1 < inC)
{
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
讀取下一個channel的數(shù)據(jù);
}
}
__syncthreads();
}
if (oc + 1 < outC)
{
讀取下一個kernel數(shù)據(jù);
}
__syncthreads();
// 注意這樣的循環(huán)順序下已經(jīng)完成了一個outC的計算
for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height; i += out_tile_row_stride;)
{
寫回當(dāng)前outC的數(shù)據(jù);
}
// 預(yù)讀取下一個in channel數(shù)據(jù) 需要注意這時候要從頭讀了
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
讀取第一個channel的數(shù)據(jù);
}
}
到這里其實我們就完成了大部分內(nèi)容了,整體骨架就是這樣,其余就是一些細(xì)節(jié)上的下標(biāo)計算問題了
3.4 一些雜項卻又需要細(xì)節(jié)
3.4.1 中間結(jié)果存儲設(shè)計
可以看到我們的偽代碼中循環(huán)順序是先oc再ic
可以想象一下,如果你先ic再oc的話,這樣確實是我們只需要遍歷一遍ic,oc多次遍歷
但是我們也要考慮寫回部分,寫回你還需要單獨再去寫,理論上先ic的話會快一些
這里就不給大家放圖了,讀者可以自己想象一下兩種計算順序的區(qū)別
需要注意的是
線程能利用的硬件資源是有限的,一個warp共用一個SM上的寄存器,具體到每個線程大概32-255個寄存器(來源于chatGPT,不嚴(yán)謹(jǐn),需要核實,后面gpt又說v100一個線程可以用800個..)
總之我們還是能少用就少用幾個
當(dāng)register存不下我們這些中間變量,就會放到local memory中
所謂的local memory是位于GMEM上的,如果發(fā)生這種情況,每次讀取中間結(jié)果
你還得跑到GMEM上去訪存,是非常之浪費時間的
兩種循環(huán)其實需要的register數(shù)目都是oc×2(2是因為你一個線程要負(fù)責(zé)好幾個位置的)
出于修正考慮,哥們兒直接開4倍,保證不會越界
3.4.2 下標(biāo)計算
這部分其實,你串行算的明白,你并行就算的明白,我們舉幾個例子來說明一下
FETCH_FLOAT4(load_reg[0]) =
FETCH_FLOAT4(in[begin_pos + OFFSET(in_tile_row_start + i, in_tile_col, inW)]);
s_in[in_tile_row_start + i][in_tile_col] = load_reg[0];
s_in[in_tile_row_start + i][in_tile_col + 1] = load_reg[1];
s_in[in_tile_row_start + i][in_tile_col + 2] = load_reg[2];
s_in[in_tile_row_start + i][in_tile_col + 3] = load_reg[3];
這里是利用向量指令去一次讀取4個32位數(shù)據(jù),s_in是開在SMEM上的,in是GMEM上的一位數(shù)據(jù)
那么可以看這個后面的下標(biāo)
begin_pos 代表當(dāng)前block的起始位序
OFFSET 是一個宏定義,代表行×一行元素數(shù)目
in[xxx] 下標(biāo)其實就是當(dāng)前block位置+block內(nèi)的位置
再看一個寫入中間結(jié)果的位置
temp_pos = i / out_tile_row_stride + j +
oc * (cur_out_block_height / out_tile_row_stride + 1);
這里要考慮到線程是在計算它負(fù)責(zé)的第幾個元素,那么就要用i / out_tile_row_stride來判斷
如果處理多個元素,那你還得用j來控制一下當(dāng)前是第幾個元素
還要考慮到不同的oc,一個oc內(nèi)負(fù)責(zé)的元素有cur_out_block_height / out_tile_row_stride +1這么多個
我們再看一個
out_pos = oc * outH * outW +
block_row * BLOCK_HEIGHT * outW + block_col * BLOCK_WIDTH +
OFFSET(out_tile_row_start + i, out_tile_col + j, outW);
首先略過幾個oc的范圍,再計算當(dāng)前block的起始位置,再計算上block內(nèi)的相對位置
每個下標(biāo)都要明白其計算的含義,本例中有很多公共表達(dá)式?jīng)]有提取出來提前計算,會影響一定性能
3.5 完整代碼
完整代碼已上傳:
https://github.com/Pegessi/conv2d_direct
3.6 性能測試
雖然是娛樂測試,但是也嚴(yán)謹(jǐn)一點,可以發(fā)現(xiàn)這個代碼會受channel數(shù)目影響很大
代碼還有一點小bug,不過不影響你執(zhí)行,大家可能會發(fā)現(xiàn)(亟待修復(fù))
不同數(shù)據(jù)規(guī)模下性能在cudnn的1/10到10倍上下橫跳,有空給大家測一下放個完整的圖。