cuda程式最佳化-2.訪存最佳化

SunStriKE發表於2024-06-17

簡介

在CUDA程式中, 訪存最佳化個人認為是最重要的最佳化項. 往往kernel會卡在資料傳輸而不是計算上, 為了最大限度利用GPU的計算能力, 我們需要根據GPU硬體架構對kernel訪存進行合理的編寫.

這章主要以計算一個tensor的模為例, 來看具體如何最佳化訪存從而提升並行效率. 以下程式碼都只列舉了kernel部分, 省略了host提交kernel的部分. Block_size均為256, 程式碼均在A100上評估效能.

\[\| \vec{v} \| = \sqrt{v_1^2 + v_2^2 + \cdots + v_n^2} \]

CPU程式碼

void cpu_tensor_norm(float* tensor, float* result, size_t len) {
    double sum = 0.0;
    for (auto i = 0; i < len; ++i) {
        sum += tensor[i] * tensor[i];
    }
    *result = sqrt(sum);
}

是一個非常簡單的函式, 當然這裡針對CPU程式也有最佳化方法, 當前實現並不是最優解. 此處忽略不表

GPU實現1-基於CPU並行思想

在傳統的多執行緒CPU任務中, 如果想處理一個超大陣列的求和, 很自然的會想到起N個執行緒, 每個執行緒算1/N的和, 最後再把這N個和加到一塊求sqrt. 根據這個思路實現了kernel1, 使用1個block, block中每個執行緒計算1/N連續儲存的資料.

__global__ void norm_kernel1(float* tensor, float* result, size_t len) {
    auto tid = threadIdx.x;
    auto size = len / blockDim.x;
    double sum = 0;

    for (auto i = tid * size; i < (tid + 1) * size; i++) {
        sum += tensor[i] * tensor[i];
    }
    result[tid] = float(sum);
}
//norm_kernel1<<<1, block_size>>>(d_tensor, d_result1, len);

image

接下來我們使用nsight-compute來分析下第一個kernel實現哪些地方不合理. 具體使用方法可以從官方文件獲取, 這裡舉個命令例子: nv-nsight-cu-cli -f --target-processes all -o profile --set full --devices 0 ./output/norm_kernel_bench
image

ncu分析給出了2個主要問題:

grid太小

因為我們只使用了256個執行緒. 而A100光CUDA core就7000個. 在GPU中不像CPU中執行緒上下文切換具有很高的成本, 如果想充分利用算力, 就要用盡量多的執行緒來提升併發度. 同時執行緒數也不能無限制增加, 因為如果每個執行緒使用了32個暫存器, 而SM中最多16384的暫存器空間的話, 多於16384/32=512個執行緒後, 這些多出來的執行緒就需要把資料存到視訊記憶體裡, 反而會降低執行效率.

讀視訊記憶體瓶頸

下面開啟detail後也給出了問題日誌: Uncoalesced global access, expected 262144 transactions, got 2097152 (8.00x) at PC

coalesced指的是視訊記憶體讀取需要是連續的, 這裡也許你會有疑問, 在kernel1裡就是按照連續的視訊記憶體讀的呀. 這裡涉及到GPU的實際執行方式, 當一個thread在等讀視訊記憶體資料完成的時候, GPU會切換到下一個thread, 也就是說是需要讓thread1讀視訊記憶體的資料和thread2的資料是連續的才會提升視訊記憶體的讀取效率, 在kernel1中明顯不連續. 同時參考ProgrammingGuide 中的描述, 每個warp對視訊記憶體的訪問是需要對齊 32/64/128 bytes(一次資料傳輸的資料量在預設情況下是32位元組), 如果所有資料傳輸處理的資料都是該warp所需要的，那麼合併度為100%，即合併訪問

GPU實現2-最佳化coalesced

__global__ void norm_kernel2(float* tensor, float* result, size_t len) {
    auto tid = threadIdx.x;
    double sum = 0.0;
    while (tid < len) {
        sum += tensor[tid] * tensor[tid];
        tid += blockDim.x;
    }
    result[threadIdx.x] = float(sum);
}
//norm_kernel2<<<1, block_size>>>(d_tensor, d_result1, len);

依然還是1個block256個執行緒, kernel2將每個thread讀取方式改成了每隔256個float讀一個, 這樣uncoalesced的報錯就不見了. 但是! 一跑bench卻發現為啥反而還變慢了呢?
image

此時可以點開memory WorkLoad Analysis

image

從這裡可以看到L2 cache命中率降低了50%左右, 原因是因為按照kernel1的訪問方式, 第一次訪問了32bytes長度, 但是隻用了一部分後, 剩下的會快取在L2中, 而Kernel2雖然訪問視訊記憶體連續了, 但每次的cache命中率會隨著讀入資料利用效率的變高而降低, 根本原因是因為執行緒和block太少導致的. 另外這張圖上還有個明顯的歎號, 我們沒有合理的用到shared_memory. 接下來的kernel3重點最佳化這兩部分

GPU實現3-增大併發&利用shared_memory

__global__ void norm_kernel3(float* tensor, float* result, size_t len) {
    auto tid = threadIdx.x + blockIdx.x * blockDim.x;
    extern __shared__ double sum[];
    auto loop_stride = gridDim.x * blockDim.x;
    sum[threadIdx.x] = 0;
    while (tid < len) {
        sum[threadIdx.x] += tensor[tid] * tensor[tid];
        tid += loop_stride;
    }
    __syncthreads();
    if (threadIdx.x == 0) {
        for (auto i = 1; i < blockDim.x; ++i) {
            sum[0] += sum[i];
        }
        result[blockIdx.x] = float(sum[0]);
    }
}
//grid_size3=64
//norm_kernel3<<<grid_size3, block_size, block_size * sizeof(double)>>>(d_tensor, d_result3, len);
//for (auto i = 1; i < grid_size3; ++i) {
//    h_result3[0] += h_result3[i];
//}
//sqrt(h_result3[0])

image

短短几行改動, 讓程式快了27倍, 這是什麼魔法(黑人問號)? kernel3做了如下幾個最佳化:

  1. 使用shared_memory用來儲存臨時加和, 最後在每個block的第一個thread裡把這些加和再加到一塊, 最後再寫回HBM. shared_memory訪問速度19T/s, HBM速度才1.5TB/s, 所以我們如果有需要臨時儲存的變數, 要儘可能的把shared_mem利用起來.

  2. 這次使用了64個block, GPU的其他SM終於不用看戲了. 但其實還是可以增加的, A100有108個SM呢, 讓我們把他用滿再看下效能auto grid_size3 = (len + block_size - 1) / block_size; 可以看到我們終於讓計算吞吐打滿了~

image

到這裡還有2個問題需要解決: 1.我們只是在GPU裡求了區域性加和, 全域性和還得挪到CPU算好挫. 2.每個block在syncthread之後只有第一個執行緒在計算, 能不能加快計算的同時減少一下計算量

GPU實現4-最佳化加和

template <int64_t BLOCK_SIZE>
__global__ void norm_kernel4(float* tensor, double* result, size_t len) {
    using BlockReduce = cub::BlockReduce<double, BLOCK_SIZE>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    double sum = 0.0;
    if (tid < len) {
        sum += tensor[tid] * tensor[tid];
    }
    double block_sum = BlockReduce(temp_storage).Sum(sum);
    if (threadIdx.x == 0) {
        atomicAdd(result, block_sum);
    }
}
//norm_kernel4<block_size><<<grid_size, block_size>>>(d_tensor, d_result4, len);
//sqrt(h_result4)

image

對與第一個執行緒連續加和的問題, 我偷懶使用了cub::BlockReduce方法, BlockReduce的原理是經典的樹形規約演算法. 利用分治的思想, 把原來的32輪加和可以簡化為5輪加和, 這樣就能極大減少長尾執行緒的計算量.

image

對於視訊記憶體上的全域性求和問題, 由於block之間是沒有任何關聯的, 我們必須使用原子操作來解決對全域性視訊記憶體的操作. 這裡為了減少原子寫衝突, 只在block的第一個執行緒上進行原子加. 另外我們可以不用cub的BlockReduce最佳化到和他差不多的效能嗎?

GPU實現5-自己實現BlockReduce

template <int64_t BLOCK_SIZE>
__global__ void norm_kernel5(float* tensor, double* result, size_t len) {
    using WarpReduce = cub::WarpReduce<double>;
    const int64_t warp_size = BLOCK_SIZE / 32;
    __shared__ typename WarpReduce::TempStorage temp_storage[warp_size];
    __shared__ float reduce_sum[BLOCK_SIZE];
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    int warp_id = threadIdx.x / 32;
    float sum = 0.0;
    if (tid < len) {
        sum += tensor[tid] * tensor[tid];
    }
    auto warp_sum = WarpReduce(temp_storage[warp_id]).Sum(sum);
    reduce_sum[threadIdx.x] = warp_sum; //這裡儘量避免wrap內的分支

    __syncthreads();
    //樹形規約
    
    int offset = warp_size >> 1;
    if (threadIdx.x % 32 == 0) {
        while (offset > 0) {
            if (warp_id < offset) {
                reduce_sum[warp_id * 32] += reduce_sum[(warp_id + offset) * 32];
            }
            offset >>= 1;
            __syncthreads();
        }
    }
    if (threadIdx.x == 0) {
        atomicAdd(result, reduce_sum[0]);
    }
}
//norm_kernel5<block_size><<<grid_size5, block_size>>>(d_tensor, d_result5, len);
//sqrt(h_result5)

image

kernel5主要分為兩個部分, 第一步進行了warp內的規約, 第二步手動實現了樹形規約的方法. 耗時基本與BlockReduce一致, 由於warp內的32個執行緒會共享暫存器和shared_mem讀寫, 在warp內先做一些規約可以適當減少後續的sync_threads輪數.
其實這個kernel還是有很大最佳化空間, 篇幅受限原因更深層次的最佳化技巧在後續說明

計算分析

image

在彙編指令統計這塊可以看到LDS(從shared_mem載入指令), LOP3(logic操作), STS(往mem中寫入操作), BAR(barrier), BSYNC(執行緒同步時的barrier, 對應atomic操作), WARPSYNC(warp同步)這些相較於kernel4多了很多. WARPSYNC/LOP3/BAR這些變多是正常的, 是kernel5裡新增的邏輯. LDS/STS增加應該是因為我們在warp_reduce時的頻率比block_reduce對共享記憶體的訪問頻率更高導致.

訪存注意要點-bank conflicts

image

記憶體分析裡出現了一個新名詞bank conflicts , 先根據官方文件 解釋下這個名詞.

SharedMemory結構: 放在shared memory中的資料是以4bytes作為1個word，依次放在32個banks中。第i個word存放在第(i % 32)個bank上。每個bank在每個cycle的 bandwidth為4bytes。所以shared memory在每個cycle的bandwidth為128bytes。這也意味著每次記憶體訪問只會訪問128bytes資料

image

如果同一個warp內的多個threads同時訪問同一個bank內的同一個word, 會觸發broadcast機制, 會同時發給多個thread. 不會產生衝突

衝突主要產生在多個threads訪問同一個bank內的不同word, 如上圖的第二列. 這樣就會導致本來的一次memory transaction被強制拆分成了2次, 而且需要序列訪問memory

解決方法:

透過錯位的方式訪問陣列, 避免訪問步長和32產生交集, 每個執行緒根據執行緒編號tid與訪問步長s的乘積來訪問陣列的32-bits字(word)：

extern __shared__ float shared[];
float data = shared[baseIndex + s * tid];

如果按照上面的方式，那麼當s*n是bank的數量(即32)的整數倍時或者說n是32/d的整數倍(d是32和s的最大公約數)時，執行緒tid和執行緒tid+n會訪問相同的bank。我們不難知道如果tid與tid+n位於同一個warp時，就會發生bank衝突，相反則不會。

仔細思考你會發現，只有warp的大小(即32)小於等於32/d時，才不會有bank衝突，而只有當d等於1時才能滿足這個條件。要想讓32和s的最大公約數d為1，s必須為奇數。於是，這裡有一個顯而易見的結論：當訪問步長s為奇數時，就不會發生bank衝突。

相關文章