高效、易用、可拓展我全都要:OneFlow CUDA Elementwise模板庫的設計優化思路

語言: CN / TW / HK

撰文|鄭澤康、姚遲、郭冉、柳俊丞

逐元素操作(也叫 Elementwise 操作)是指對 Tensor 中的每個元素應用一個函式變換,得到最終輸出結果。在深度學習裡,有很多運算元屬於 Elementwise 運算元範疇,比如常用的啟用函式(如ReLU、GELU ),ScalarMultiply(對 Tensor 每個元素都乘上一個標量)等操作。

為此,OneFlow 針對這種 Elementwise 操作抽象出一套 CUDA 模板, 開發者只需把計算邏輯封裝到一個結構體內,即可獲得一個 CUDA Elementwise 運算元 ,以 ReLU 為例:

// Write ReLU Functor. 
template<typename T>
struct ReluFunctor {
  OF_DEVICE_FUNC T operator()(T x) const {
    const T zero_val = static_cast<T>(0); 
    return (x > zero_val) ? x : zero_val; 
  }
};

// Use CUDA Elementwise Template. 
OF_CUDA_CHECK((cuda::elementwise::Unary(ReluFunctor<T>(), elem_cnt, dx->mut_dptr<T>(),
                                        x->dptr<T>(), ctx->stream()->As<ep::CudaStream>()->cuda_stream())));

這樣一套簡單易用的 Elementwise 模板 不僅提高了開發效率,也能保證計算效能 我們在 NVIDIA A100 40GB 環境下使用 Nsight Compute,和 PyTorch 的 Cast 運算元進行測試,測試用例是將 float32 型別的 Tensor 轉換為 half 型別,比較兩者的執行時間和頻寬,在各個資料形狀情況下, OneFlow 均能比 PyTorch快 80-90%,並接近機器理論頻寬。

下面我們會逐個介紹這套模板的設計思路以及優化技巧。

1

設定合理的 BlockSize 和 GridSize

關於設定執行緒塊個數和執行緒數量的問題,我們在《 如何設定CUDA Kernel中的grid_size和block_size? 》一文中有討論過,這裡我們的設定規則還稍微有點區別。在CUDA 官方文件 Compute Capabilities( https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities )中提到了:

  • 主流架構裡,每個 Block 最大暫存器數量是 64 K

  • 每個執行緒所能使用的最大暫存器數量是 255 個

在使用最大暫存器數量的前提下,那每個 Block 最多能啟動  64 * 1024 / 255 = 256 個執行緒(往2的倍數取整),因此這裡我們設定了一個常量  constexpr int kBlockSize = 256; 而 Grid Size 大小的設定規則在  GetNumBlocks  這個函式中:

constexpr int kBlockSize = 256
constexpr int kNumWaves = 32;

inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
  ...
  /*   n: The number of the elements.    sm_count: The number of the SM.    tpm: The maximum resident threads in per multiprocessor.    */
  *num_blocks = std::max<int>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
                                                   sm_count * tpm / kBlockSize * kNumWaves));
  return cudaSuccess;
}
  • 執行緒塊最小個數為1

  • 執行緒塊最大個數是從  處理所有元素所需最小的執行緒總數 和  wave 數目*GPU 一次可以排程 SM 數量 * 每個 SM 最大 block 數 中取最小值,這裡我們的 wave 數目設定為固定32大小

在資料量較小的情況下,不會啟動過多的執行緒塊。在資料量較大的情況下,儘可能將執行緒塊數目設定為數量足夠多的整數個 wave,以保證 GPU 實際利用率夠高。

2

使用向量化操作

大部分 Elementwise 運算元的計算邏輯較為簡單,瓶頸主要是在頻寬利用上。英偉達的部落格CUDA Pro Tip: Increase Performance with Vectorized Memory Access( https://developer.nvidia.com/blog/cuda-pro-tip-increase-performance-with-vectorized-memory-access/ )提到,使用向量化操作能夠提升讀寫的頻寬,而 CUDA 裡也提供了一系列資料型別來支援向量化操作,如 float2float4 ,就是將2個或4個 float 資料作為一個整體。在一些高效能訓練推理庫如  LightSeq  就使用了大量的  float4  型別:

template <typename T>
__global__ void ker_layer_norm(T *ln_res, T *vars, T *means, const T *inp,                                const T *scale, const T *bias, int hidden_size) {
  // step 0. compute local sum
  float l_sum = 0;
  float l_square_sum = 0;
  const float4 *inp_f4 = (const float4 *)inp + blockIdx.x * hidden_size; // use float4
  for (uint idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
    float4 val = inp_f4[idx];
    ...
  }
}

在實際中,我們的運算元需要支援不同資料型別(如 int, half ),如果採用 CUDA 內建的向量化資料型別操作,顯然要給每個運算元寫多個版本,增加了開發負擔。為此我們實現了一個  Pack 資料結構,用於靈活支援不同資料型別的向量化。

我們先定義了一個  PackType 型別型別來代表向量化的資料,它代表的(向量化後的)資料大小為  sizeof(T) * pack_size

template<typename T, int pack_size>
struct GetPackType {
  using type = typename std::aligned_storage<pack_size * sizeof(T), pack_size * sizeof(T)>::type;
};

template<typename T, int pack_size>
using PackType = typename GetPackType<T, pack_size>::type;

然後實現了一個  union 型別  Pack ,它內部定義了  PackType<T, pack_size> storage; 來佔用空間:

template<typename T, int pack_size>
union Pack {
  static_assert(sizeof(PackType<T, pack_size>) == sizeof(T) * pack_size, "");
  __device__ Pack() {
    // do nothing
  }
  PackType<T, pack_size> storage;
  T elem[pack_size];
};

與  storage 共享記憶體的,還有  T elem[pack_size];這樣方便後續的 Elementwise 操作:在後續計算裡,我們對 elem 陣列中的每個元素都應用 functor ,得到輸出結果。

CUDA 裡最大支援128 bit 的 pack 大小,而在浮點資料型別中,最小的型別(half)大小為16 bit,最多能把128 / 16=8 個 half 資料 pack 到一起,因此我們設定了這兩個常量, kMaxPackBytes 表示 pack 最大位元組數, kMaxPackSize 表示 pack 資料的最大個數:

constexpr int kMaxPackBytes = 128 / 8;
constexpr int kMaxPackSize = 8;

3

呼叫鏈

跟蹤  oneflow/core/cuda/elementwise.cuh 中的實現,會發現,這套模板會分別為一元、二元、三元的 Elementwise 提供介面: UnaryBinaryTernary ,文章開始處的  ReLU 運算元就使用了  Unary 的介面。進一步分析可以發現,它們經過層層呼叫後,其實最終都會呼叫到  ApplyGeneric ,基本呼叫關係如下:

Unary/Binary/Ternary
  -> xxxFactory
     -> GenericLauncher<...>::Launch
       -> ApplyGeneric(CUDA Kernel)

ApplyGeneric 這個 CUDA Kernel 中所做的主要工作是:

  • 根據引數建立一個  functor

  • 進入迴圈,針對打包(pack)後的資料,呼叫  ApplyPack 函式,每呼叫一次  ApplyPack ,就處理一批 pack 後的資料

  • 當最後存在元素個數不能被  pack_size 整除的情況時,需要讓執行緒處理下尾部剩餘元素

實現程式碼如下:

template<int pack_size, bool tail, typename FactoryT, typename R, typename... IN>
__global__ void __launch_bounds__(kBlockSize)
    ApplyGeneric(FactoryT factory, int64_t n_pack, PackType<R, pack_size>* pack_r,
                 const PackType<IN, pack_size>*... pack_in, int64_t n_tail, R* tail_r,
                 const IN*... tail_in) {
  auto functor = factory();
  const int global_tid = blockIdx.x * kBlockSize + threadIdx.x;
  for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) {
    pack_r[i] = ApplyPack<pack_size, decltype(functor), R, IN...>(
        functor, (FetchPack<IN, pack_size>(pack_in + i).elem)...);
  }
  if (tail && global_tid < n_tail) { tail_r[global_tid] = functor((tail_in[global_tid])...); }
}

ApplyPack 函式定義如下,它對一個  pack 內的元素做了個迴圈,對  elem 陣列中的每個元素呼叫  functor ,得到輸出結果並返回:

template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__
    typename std::enable_if<HasApply2<FunctorT>::value == false, PackType<R, pack_size>>::type
    ApplyPack(const FunctorT& functor, const IN... in[pack_size]) {
  Pack<R, pack_size> ret;
#pragma unroll
  for (int j = 0; j < pack_size; ++j) { ret.elem[j] = functor((in[j])...); }
  return ret.storage;
}

整個  Elementwise 運算元呼叫流程如下所示:

4

針對 half2 資料型別優化

在 half 資料型別下,如果直接對其進行操作,其運算元頻寬是跟 float32 型別相當的。CUDA 官方有針對 half2 推出一系列特殊指令,如  hadd2 就可以實現兩個  half2 資料的加法,進而提高吞吐量。

考慮到這種情況,OneFlow 給  ApplyPack 函式特化了一個版本,通過呼叫 functor 的  apply2 函式,來呼叫 half2 相關特殊指令,介面如下:

template<int pack_size, typename FunctorT, typename R, typename... IN>
__device__ typename std::enable_if<HasApply2<FunctorT>::value == true && pack_size % 2 == 0,
                                   PackType<R, pack_size>>::type
ApplyPack(const FunctorT& functor, const IN... in[pack_size]) {
  Pack<R, pack_size> ret;
#pragma unroll
  for (int j = 0; j < pack_size; j += 2) { functor.Apply2(ret.elem + j, (in + j)...); }
  return ret.storage;
}

以先前的  Cast 運算元為例,我們在  CastFunctor 內部通過呼叫  __float22half2_rn 指令,將一個 float2 資料轉換為一個 half2 資料。

template<typename From>
struct CastFunctor<half, From, typename std::enable_if<!std::is_same<From, half>::value>::type> {
  ...

  __device__ void Apply2(half* to, const From* from) const {
    float2 f2;
    f2.x = static_cast<float>(from[0]);
    f2.y = static_cast<float>(from[1]);
    *reinterpret_cast<half2*>(to) = __float22half2_rn(f2);
  }
};

5

擴充套件多元操作

前面已經提到,現有的 OneFlow 模板,將 Elementwise 運算元進一步分為一元、二元、三元操作。並利用工廠模式,使得他們最終統一呼叫  ApplyGeneric 。這種設計方式易於拓展:當需要支援更多輸入的操作時,只需要編寫對應的工廠即可。

template<typename FunctorT>
struct SimpleFactory {
  explicit SimpleFactory(FunctorT functor) : tpl(functor) {}
  __device__ FunctorT operator()() const { return tpl; }

 private:
  FunctorT tpl;
};

template<typename FactoryT, typename R, typename A>
inline cudaError_t UnaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a,
                                    cudaStream_t stream) {
  return GenericLauncher<FactoryT, R, A>::Launch(factory, n, r, a, stream);
}

template<typename FunctorT, typename R, typename A>
inline cudaError_t Unary(FunctorT functor, int64_t n, R* r, const A* a, cudaStream_t stream) {
  return UnaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, stream);
}

// BinaryWithFactory TernaryWithFactory ... 
// Binary Ternary ...

至此,OneFlow 的高效能 CUDA Elementwise 模板的設計,優化手段就介紹完畢,最後再來總結下這套模板的優勢:

  • 效能夠高,應用這套 Elementwise 模板的運算元都能打滿機器的頻寬,速度也夠快。

  • 開發效率高,開發人員可以不用過分關注 CUDA 邏輯及相關優化手段,只需要編寫計算邏輯即可。

  • 可擴充套件性強,目前這套模板支援了一元,二元,三元操作。若今後有需求拓展,支援更多輸入時,只需要仿照編寫對應的工廠即可。

其他人都在看

點選“ 閱讀原文 ,歡迎下載體驗OneFlow新一代開源深度學習框架