gemv優(yōu)化總結(jié)

有朋友來(lái)信說(shuō)：

1. “除了以NVIDIA（英偉達(dá)）為例，能不能談點(diǎn)國(guó)產(chǎn)GPU優(yōu)化的經(jīng)驗(yàn)分享？”

2. “老講國(guó)外的東西，你們能不能支持一下國(guó)產(chǎn)CPU和加速卡？“

這里解釋一下原因：

1. N卡的資料和環(huán)境大家都比較好找，對(duì)于學(xué)習(xí)GPU并行優(yōu)化編程的朋友比較友善。

2. 暫時(shí)受限于商業(yè)保密，我們相信后續(xù)會(huì)逐步開(kāi)放起來(lái)，學(xué)習(xí)的平臺(tái)和環(huán)境也容易找到。到時(shí)就可以分享一些國(guó)產(chǎn)CPU和加速卡的優(yōu)化經(jīng)驗(yàn)出來(lái)。

------ 正文分割線 ------

本文主要是介紹如何對(duì)gemv算法進(jìn)行優(yōu)化。gemv，即矩陣向量乘，即計(jì)算一個(gè)矩陣A與一個(gè)向量x的乘積，這是并行計(jì)算中的經(jīng)典話題。個(gè)人感覺(jué)，gemv的優(yōu)化核心是需要考慮不同shape的情況，然后針對(duì)型地進(jìn)行優(yōu)化。本篇文章會(huì)先介紹一下針對(duì)不同shape設(shè)計(jì)不同的并行算法，然后說(shuō)明一下優(yōu)化思路和相關(guān)優(yōu)化技巧，最后說(shuō)一下實(shí)驗(yàn)效果，在A矩陣列數(shù)為16 128的時(shí)候，我寫(xiě)的gemv能擁有超越cublas的性能表現(xiàn)。

一、前言

首先介紹一下gemv算法。給定矩陣A和向量x，gemv需要計(jì)算兩者的乘積，示意圖如下：

gemv

二、針對(duì)不同shape的并行算法設(shè)計(jì)

這次講到并行算法設(shè)計(jì)，什么叫并行算法設(shè)計(jì)。每個(gè)人的理解都不太一樣，在GPU中，我的理解就是：設(shè)計(jì)block和thread的workload，說(shuō)白了就是要搞清楚一個(gè)block負(fù)責(zé)哪部分的計(jì)算，一個(gè)thread要負(fù)責(zé)哪部分的計(jì)算。而設(shè)計(jì)的原則就是盡可能地減少訪存，提高數(shù)據(jù)的復(fù)用概率，然后讓所有的處理器都滿負(fù)荷地進(jìn)行工作，不能浪費(fèi)。

2.1 針對(duì)n=32

對(duì)于n=32的情況，我們將每個(gè)block設(shè)置為256個(gè)線程，4個(gè)warp，然后每個(gè)warp負(fù)責(zé)一行元素的計(jì)算。每個(gè)warp要對(duì)x進(jìn)行訪問(wèn)，然后在warp內(nèi)部進(jìn)行一次reduce求和操作。

n=32

代碼如下：

template 
__device__ __forceinline__ float warpReduceSum(float sum) {
    if (WarpSize >= 32)sum += __shfl_down_sync(0xffffffff, sum, 16); // 0-16, 1-17, 2-18, etc.
    if (WarpSize >= 16)sum += __shfl_down_sync(0xffffffff, sum, 8);// 0-8, 1-9, 2-10, etc.
    if (WarpSize >= 8)sum += __shfl_down_sync(0xffffffff, sum, 4);// 0-4, 1-5, 2-6, etc.
    if (WarpSize >= 4)sum += __shfl_down_sync(0xffffffff, sum, 2);// 0-2, 1-3, 4-6, 5-7, etc.
    if (WarpSize >= 2)sum += __shfl_down_sync(0xffffffff, sum, 1);// 0-1, 2-3, 4-5, etc.
    return sum;
}

// if N == 32
__global__ void Sgemv_v0( 
    float * __restrict__ A,
    float * __restrict__ x,
    float * __restrict__ y, 
    const int M,
    const int N) {
    // Block index
    int bx = blockIdx.x;

    // Thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    const int warp_size=32;
    int laneId= tx % warp_size;
    int current_row = blockDim.y * bx + ty;

    if(current_row < M){
        float res=0;
        int kIteration = N/warp_size;
        if(kIteration==0) kIteration=1;
        #pragma unroll
        for(int i=0; i< kIteration; i++){
            int current_col = i*warp_size + laneId;
            res += A[current_row*N + current_col] * x[current_col];
        }
        res = warpReduceSum(res);
        if(laneId==0) y[current_row]=res;
    }
}

2.2 針對(duì)n=128

對(duì)于n=128的情況，同樣讓warp負(fù)責(zé)一行元素的計(jì)算，但是因?yàn)槊啃械脑乇容^多，所以采用了float4進(jìn)行向量化的訪存。能夠有更高的訪存效率。

n=128

代碼如下：

template 
__device__ __forceinline__ float warpReduceSum(float sum) {
    if (WarpSize >= 32)sum += __shfl_down_sync(0xffffffff, sum, 16); // 0-16, 1-17, 2-18, etc.
    if (WarpSize >= 16)sum += __shfl_down_sync(0xffffffff, sum, 8);// 0-8, 1-9, 2-10, etc.
    if (WarpSize >= 8)sum += __shfl_down_sync(0xffffffff, sum, 4);// 0-4, 1-5, 2-6, etc.
    if (WarpSize >= 4)sum += __shfl_down_sync(0xffffffff, sum, 2);// 0-2, 1-3, 4-6, 5-7, etc.
    if (WarpSize >= 2)sum += __shfl_down_sync(0xffffffff, sum, 1);// 0-1, 2-3, 4-5, etc.
    return sum;
}

// if N>= 128
__global__ void Sgemv_v1( 
    float * __restrict__ A,
    float * __restrict__ x,
    float * __restrict__ y, 
    const int M,
    const int N) {
    // Block index
    int bx = blockIdx.x;

    // Thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    const int warp_size=32;
    int laneId= tx % warp_size;
    int current_row = blockDim.y * bx + ty;

    if(current_row < M){
        float res=0;
        int kIteration = (N/warp_size)/4;
        if(kIteration==0) kIteration=1;
        A = &A[current_row*N];
        #pragma unroll
        for(int i=0; i< kIteration; i++){
            int current_col_vec = (i*warp_size + laneId);
            float4 current_val= reinterpret_cast(A)[current_col_vec];
            float4 current_x = reinterpret_cast(x)[current_col_vec];
            res += current_val.x*current_x.x;
            res += current_val.y*current_x.y;
            res += current_val.z*current_x.z;
            res += current_val.w*current_x.w;
        }
        res = warpReduceSum(res);
        if(laneId==0) y[current_row]=res;
    }
}

2.3 針對(duì)n=16

對(duì)于n=16的情況，讓一個(gè)warp負(fù)責(zé)兩行元素的計(jì)算。以warp0為例，0-15號(hào)線程負(fù)責(zé)第0行元素的計(jì)算，而16-31號(hào)線程負(fù)責(zé)第1行元素的計(jì)算。

n=16

代碼如下：

template 
__device__ __forceinline__ float warpReduceSum(float sum) {
    if (WarpSize >= 32)sum += __shfl_down_sync(0xffffffff, sum, 16); // 0-16, 1-17, 2-18, etc.
    if (WarpSize >= 16)sum += __shfl_down_sync(0xffffffff, sum, 8);// 0-8, 1-9, 2-10, etc.
    if (WarpSize >= 8)sum += __shfl_down_sync(0xffffffff, sum, 4);// 0-4, 1-5, 2-6, etc.
    if (WarpSize >= 4)sum += __shfl_down_sync(0xffffffff, sum, 2);// 0-2, 1-3, 4-6, 5-7, etc.
    if (WarpSize >= 2)sum += __shfl_down_sync(0xffffffff, sum, 1);// 0-1, 2-3, 4-5, etc.
    return sum;
}

// if N <= 16
template <
    const int ROW_PER_WARP
    > 
__global__ void Sgemv_v2( 
    float * __restrict__ A,
    float * __restrict__ x,
    float * __restrict__ y, 
    const int M,
    const int N) {
    // Block index
    int bx = blockIdx.x;

    // Thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    const int warp_size=32;
    int laneId= tx % warp_size;
    int current_warp_row = (blockDim.y * bx + ty) * ROW_PER_WARP;
    const int kWarp_size = warp_size / ROW_PER_WARP;
    int kLaneId = laneId % kWarp_size;
    int current_thread_row = current_warp_row + laneId / kWarp_size;

    if(current_thread_row < M){
        float res=0;
        int current_col = kLaneId;
        res += A[current_thread_row * N + current_col] * x[current_col];
        res = warpReduceSum(res);
        if(kLaneId==0) y[current_thread_row]=res;
    }
}

三、優(yōu)化思路：

上一節(jié)說(shuō)明了如何針對(duì)不同維度的n進(jìn)行優(yōu)化，這一節(jié)說(shuō)明一下為什么要這么設(shè)計(jì)，以及這樣的設(shè)計(jì)方式能夠帶來(lái)什么樣的好處。主要考慮的因素有兩個(gè)，如下：

3.1 盡可能地讓warp中的32個(gè)線程忙碌

這個(gè)主要是針對(duì)n<32的情況，例如n=16，如果使用一個(gè)warp來(lái)負(fù)責(zé)一行元素的計(jì)算，那么warp中有一半的元素都是浪費(fèi)的。所以讓一個(gè)warp來(lái)負(fù)責(zé)多行元素的計(jì)算，這樣讓32個(gè)線程全部忙碌起來(lái)。

3.2 盡可能地提高訪存效率

① global mem->register

將數(shù)據(jù)從global memory搬運(yùn)到寄存器上時(shí)，最重要的就是考慮是不是進(jìn)行了合并訪存。在這里，我們只考慮矩陣數(shù)據(jù)在global mem中是地址對(duì)齊的，即n是2的多次冪。上述的三種并行實(shí)現(xiàn)中，warp中的32個(gè)線程都是連續(xù)地訪問(wèn)32個(gè)float或者128個(gè)float，因而滿足了合并訪存的條件，確保了global -> register的訪存效率。

② shared mem->register

說(shuō)到這里，可能會(huì)有讀者好奇，上述的代碼都沒(méi)有用到shared mem。為啥要說(shuō)這個(gè)點(diǎn)。我們可以再仔細(xì)看看上述的三種并行實(shí)現(xiàn)，以第2種為例，一個(gè)block中有4個(gè)warp，每個(gè)warp都需要對(duì)x進(jìn)行一次global上的訪存，所以一個(gè)block有4次訪存。如果將x存儲(chǔ)到shared mem中，4個(gè)warp都去訪問(wèn)shared mem上的x，這樣的話，對(duì)于global的訪存就從4次變成1次。直觀上會(huì)有性能提升，但不幸的是，如果用shared mem的話，將global mem的數(shù)據(jù)搬運(yùn)至shared mem需要有同步操作，這又會(huì)導(dǎo)致性能的下降?？偟膩?lái)說(shuō)，使用shared mem并沒(méi)有得到顯著的提升，不過(guò)還是在這里說(shuō)明一下。

③ 向量化訪存

向量化訪存就是一個(gè)老生常談的話題了，說(shuō)白了就是盡可能地使用128bit的訪存指令，這個(gè)在reduce、sgemm、elementwise專題上說(shuō)了很多，就不再多說(shuō)。

四、實(shí)驗(yàn)與總結(jié)

筆者在V100上進(jìn)行了實(shí)驗(yàn)，迭代1000次，用nsight進(jìn)行了測(cè)試，性能數(shù)據(jù)如下：

可以看出，在n=16以及n=128的情況下，都比cublas性能要好。n=32的情況要差于cublas。如果再加上向量化訪存應(yīng)該能夠有更好的性能表現(xiàn)。由于我實(shí)在沒(méi)時(shí)間再進(jìn)行深入，有心的同學(xué)可以改改代碼看看效果 :)。