CUDA優化之PReLU效能調優

語言: CN / TW / HK
時間 2022-05-23 01:02:31

 

撰文｜鄭澤康

 

InsightFace模型裡大量使用了PReLU啟用函式，而PReLU的工作模式有兩種：

 

1. PReLU(1)，此時權重alpha的形狀為(1, )，等價於一個Elementwise操作。

2. PReLU(channels)，此時權重alpha的形狀為(channels, )，和輸入特徵(N, C, H, W)中C的大小是對應的。此時PReLU等價於一個Binary Broadcast操作。

 

InsightFace模型裡的PReLU工作模式是第二種，之前已經介紹過CUDA Elementwise操作優化，而在Broadcast情形下也存在一定的優化機會。

 

1

樸素實現

 

一個樸素實現的思想就是在迴圈內部，根據當前元素的索引，推算出該元素對應需要使用的alpha權重的索引。然後判斷當前元素x是否大於0，若大於0則返回x，小於0則返回alpha*x。對應程式碼如下：

 

template<typename T>__global__ void PReluForwardGpu(const int32_t elem_cnt, const int32_t alpha_size,                                const int32_t inner_size, const T* x, const T* alpha, T* y) {  CUDA_1D_KERNEL_LOOP(i, elem_cnt) {    const T x_i = x[i];    const T alpha_i = alpha[(i / inner_size) % alpha_size];    y[i] = x_i > 0 ? x_i : x_i * alpha_i;  }}

 

其中：

 

  • inner_size表示的是通道維後面維度乘積，以NCHW格式為例，inner_size=H*W

  •  alpha_size表示通道維大小

 

在CUDA中，整數除法的計算代價是比較昂貴的（http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#maximize-instruction-throughput）關於計算指令耗時這一章中有提到：

 

Integer division and modulo operation are costly as they compile to up to 20 instructions.

 

整數除法，取餘操作會被編譯成多達20條指令。而我們這裡計算alpha的索引的時候，分別用到一次除法，一次取餘，佔整個Kernel的主要計算量，下面我們將用向量化的思路來提高讀寫頻寬的同時，減少整數除法，取餘的計算次數。

 

2

Pack向量化優化

 

我們考慮一個比較簡單的例子，輸入為(1, 2, 4, 4)，對應PReLU(2)

 

 

顯然，輸入在hw維上是連續的，在 inner_size 滿足被pack整除的條件下，一個pack內的元素應用到的是同一個alpha權重**。參見下圖：

 

 

 

這樣我們就能以向量化形式去處理元素，以提升讀寫頻寬。並且每一個pack內部只需要計算一次，向量化處理相比逐元素計算能節省不小計算量。對應程式碼如下：

 

template<typename T, typename IndexType, int pack_size>__global__ void PReluForwardMultiAlphaGpu(const IndexType elem_cnt, const IndexType alpha_size,                                          const IndexType inner_size, const T* x, const T* alpha, T* y) {  int32_t global_thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  using LoadType = cuda::elementwise::PackType<T, pack_size>;  using LoadPack = cuda::elementwise::Pack<T, pack_size>;  T zero_val = static_cast<T>(0);  for (int64_t linear_index = global_thread_id * pack_size; linear_index < elem_cnt;       linear_index += gridDim.x * blockDim.x * pack_size) {    // 計算當前Pack所使用到Alpha的索引        IndexType alpha_idx = (linear_index/inner_size%alpha_size);    const LoadType* x_load = reinterpret_cast<const LoadType*>(x + linear_index);    // 以向量化的形式載入輸入x    LoadPack x_vec;    x_vec.storage = *x_load;
    LoadPack y_vec;    // 迴圈展開，逐個處理Pack內的元素#pragma unroll    for (int i = 0; i < pack_size; i++) {      y_vec.elem[i] = x_vec.elem[i] > zero_val ? x_vec.elem[i] : x_vec.elem[i] * alpha[alpha_idx];    }    // 以向量化的形式儲存輸出y    *(reinterpret_cast<LoadType*>(y + linear_index)) = y_vec.storage;  }}

 

我們在Nsight Compute內簡單比較下優化前後的結果，測試資料為(96, 64, 112, 112)，機器為A100-40GB。藍色一欄是使用向量化優化過的kernel，而綠色一欄是樸素實現的kernel。可以看到，經過優化後，我們計算佔比降低20%-30%，吞吐提升了30+%。優化後的kernel頻寬能達到1350GB/s，已經很接近A100上的理論頻寬1555GB/s。

 

 

當然也不是所有形狀都支援向量化操作，當inner_size無法被對應的pack_size 整除時，只能退回到樸素實現上。

 

3

基準測試

 

在A100-40GB測試機器上，我們對Insightface涉及到的Tensor形狀，與PyTorch實現進行比較，測試資料如下：

 

 

經過優化PReLU的OneFlow，在大部分情況下均有比PyTorch接近2倍的領先優勢，在最後一種情況由於形狀較為特殊，無法應用向量化的優化，所以表現與PyTorch持平。

 

其他人都在看

歡迎下載體驗OneFlow v0.7.0：http://github.com/Oneflow-Inc/oneflow/

 

本文分享自微信公眾號 - OneFlow（OneFlowTechnology）。
如有侵權，請聯絡 [email protected] 刪除。
本文參與“OSC源創計劃”，歡迎正在閱讀的你也加入，一起分享。

「」
「」