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

有朋友來信說：

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

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

這里解釋一下原因：

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

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

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

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

一、前言

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

gemv

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

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

2.1 針對n=32

對于n=32的情況，我們將每個block設(shè)置為256個線程，4個warp，然后每個warp負(fù)責(zé)一行元素的計算。每個warp要對x進(jì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 針對n=128

對于n=128的情況，同樣讓warp負(fù)責(zé)一行元素的計算，但是因為每行的元素比較多，所以采用了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 針對n=16

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

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

3.1 盡可能地讓warp中的32個線程忙碌

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

3.2 盡可能地提高訪存效率

① global mem->register

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

② shared mem->register

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

③ 向量化訪存

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

四、實驗與總結(jié)

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

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