有朋友來信說：

1. “除了以NVIDIA（英偉達）為例，能不能談點國產GPU優化的經驗分享？”

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

這里解釋一下原因：

1. N卡的資料和環境大家都比較好找，對于學習GPU并行優化編程的朋友比較友善。

2. 暫時受限于商業保密，我們相信后續會逐步開放起來，學習的平臺和環境也容易找到。到時就可以分享一些國產CPU和加速卡的優化經驗出來。

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

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

一、前言

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

gemv

二、針對不同shape的并行算法設計

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

2.1 針對n=32

對于n=32的情況，我們將每個block設置為256個線程，4個warp，然后每個warp負責一行元素的計算。每個warp要對x進行訪問，然后在warp內部進行一次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負責一行元素的計算，但是因為每行的元素比較多，所以采用了float4進行向量化的訪存。能夠有更高的訪存效率。

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負責兩行元素的計算。以warp0為例，0-15號線程負責第0行元素的計算，而16-31號線程負責第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;
    }
}

三、優化思路：

上一節說明了如何針對不同維度的n進行優化，這一節說明一下為什么要這么設計，以及這樣的設計方式能夠帶來什么樣的好處。主要考慮的因素有兩個，如下：

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

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

3.2 盡可能地提高訪存效率

① global mem->register

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

② shared mem->register

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

③ 向量化訪存

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

四、實驗與總結

筆者在V100上進行了實驗，迭代1000次，用nsight進行了測試，性能數據如下：

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