0x0. 前言

如題所述，本篇文章推薦和講解一下OneFlow ElementWise模板，F(xiàn)astAtomicAdd，OneFlow UpsampleNearest2d模板的用法以及原理。但OneFlow ElementWise模板的用法和原理在【BBuf的CUDA筆記】一，解析OneFlow Element-Wise 算子實(shí)現(xiàn) 已經(jīng)講過了，所以這篇文章里不再贅述，主要講解后面2個(gè)。我將上述三個(gè)算法的實(shí)現(xiàn)都分別抽出來放到了  https://github.com/BBuf/how-to-optim-algorithm-in-cuda 這個(gè)工程的 elementwise/FastAtomicAdd/UpsampleNearest2D 三個(gè)文件夾中，并且三個(gè)算法的實(shí)現(xiàn)都分別只用一個(gè).cu文件進(jìn)行整理，使用nvcc編譯可以使用，有需要的同學(xué)請(qǐng)自取。

0x1. OneFlow elementwise模板

將 oneflow 的 elementwise 模板抽出來方便大家使用，這個(gè) elementwise 模板實(shí)現(xiàn)了高效的性能和帶寬利用率，并且用法非常靈活。完整實(shí)驗(yàn)代碼見 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/elementwise/elementwise.cu，原理講解請(qǐng)看：【BBuf 的CUDA筆記】一，解析OneFlow Element-Wise 算子實(shí)現(xiàn) 。這里以逐點(diǎn)乘（z = x * y，其中x，y，z是形狀完全一樣的Tensor）為例，性能和帶寬的測(cè)試情況如下 (A100 PCIE 40G)：

可以看到無論是性能還是帶寬，使用 oneflow 的 elementwise 模板相比于原始實(shí)現(xiàn)都有較大提升。

涉及到的主要優(yōu)化技術(shù)有向量化數(shù)據(jù)訪問，選取合適的GridSize和BlockSize，循環(huán)展開和Grid-Stride Loops等技巧。

模板代碼和用法詳見：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/elementwise/elementwise.cu

0x2. FastAtomicAdd

眾所周知，atomicAdd是CUDA中非常昂貴的操作，特別是對(duì)于half類型來說 atomicAdd 巨慢無比，慢到如果一個(gè)算法需要用到 atomicAdd，那么相比于用 half ，轉(zhuǎn)成 float ，再 atomicAdd，再轉(zhuǎn)回去還要慢很多。但是我們有時(shí)候不得不去執(zhí)行half類型的原子加，這個(gè)時(shí)候怎么能提升性能呢？

PyTorch給出了一個(gè)快速原子加的實(shí)現(xiàn)（我這里魔改了一下，去掉了一些不需要的參數(shù)，完整測(cè)試代碼見 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/fast_atomic_add_half.cu ）：

// FastAdd is referenced from
// https://github.com/pytorch/pytorch/blob/396c3b1d88d7624938a2bb0b287f2a19f1e89bb4/aten/src/ATen/native/cuda/KernelUtils.cuh#L29
template<typename T, typename std::enable_if<std::is_same<half, T>::value>::type* = nullptr>
__device__ __forceinline__ void FastSpecializedAtomicAdd(T* base, size_t offset,
                                                         const size_t length, T value) {
#if ((defined(CUDA_VERSION) && (CUDA_VERSION < 10000)) \
     || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)))
  atomicAdd(reinterpret_cast<half*>(base) + offset, static_cast<half>(value));
#else
  // Accounts for the chance base falls on an odd 16 bit alignment (ie, not 32 bit aligned)
  __half* target_addr = reinterpret_cast<__half*>(base + offset);
  bool low_byte = (reinterpret_cast<std::uintptr_t>(target_addr) % sizeof(__half2) == 0);

  if (low_byte && offset < (length - 1)) {
    __half2 value2;
    value2.x = value;
    value2.y = __float2half_rz(0);
    atomicAdd(reinterpret_cast<__half2*>(target_addr), value2);

  } else if (!low_byte && offset > 0) {
    __half2 value2;
    value2.x = __float2half_rz(0);
    value2.y = value;
    atomicAdd(reinterpret_cast<__half2*>(target_addr - 1), value2);

  } else {
    atomicAdd(reinterpret_cast<__half*>(base) + offset, static_cast<__half>(value));
  }
#endif
}

template<typename T, typename std::enable_if<!std::is_same<half, T>::value>::type* = nullptr>
__device__ __forceinline__ void FastSpecializedAtomicAdd(T* base, size_t offset,
                                                         const size_t length, T value) {
  atomicAdd(base + offset, value);
}

template<class T>
__device__ __forceinline__ void FastAdd(T* base, size_t offset, const size_t length, T value) {
  FastSpecializedAtomicAdd(base, offset, length, value);
}

也就是把half類型的原子加轉(zhuǎn)換成half2類型的原子加，為了驗(yàn)證這個(gè)快速原子加相比于half類型的原子加以及pack 2個(gè)half 到 half2再執(zhí)行原子加的性能表現(xiàn)，我實(shí)現(xiàn)了三個(gè)算法（.cu文件）。它們都是針對(duì)half數(shù)據(jù)類型做向量的內(nèi)積，都用到了atomicAdd，保證數(shù)據(jù)的長(zhǎng)度以及gridsize和blocksize都是完全一致的。具體如下：

1. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/atomic_add_half.cu 純half類型的atomicAdd。
2. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/atomic_add_half_pack2.cu half+pack，最終使用的是half2類型的atomicAdd。
3. https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/fast_atomic_add_half.cu 快速原子加，雖然沒有顯示的pack，但本質(zhì)上也是通過對(duì)單個(gè)half補(bǔ)0使用上了half2的原子加。

下面展示3個(gè)腳本通過ncu profile之后的性能表現(xiàn)：

可以看到使用pack half的方式和直接使用half的fastAtomicAdd方式得到的性能結(jié)果一致，均比原始的half的原子加快3-4倍。

接下來驗(yàn)證一下是否存在warp分支分化問題，對(duì)比了一下fastAtomicAdd和pack half2的ncu匯編代碼，并未發(fā)現(xiàn)不同類型的指令：

fastAtomicAdd 計(jì)算部分：

atomicAddhalfpack2計(jì)算部分：

每一種指令的類型都能在兩份代碼中找到，初步判斷不會(huì)因?yàn)閒astAtomicAdd實(shí)現(xiàn)中的下述if語句存在線程分化問題。

綜上所述，使用FastAtomicAdd可以大幅度提升half數(shù)據(jù)類型原子加的性能并且不需要手動(dòng)Pack，使用方法更加簡(jiǎn)單。

模板代碼和用法詳見：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/FastAtomicAdd/fast_atomic_add_half.cu

0x3. Oneflow Upsample模板

在Stable Diffusion的反向擴(kuò)散過程中使用到了UNet，而UNet中存在大量的UpsampleNearest2D上采樣。PyTorch對(duì)于UpsampleNearest都是通用的實(shí)現(xiàn)（https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/UpSampleNearest2d.cu#L112-L163） ，這種實(shí)現(xiàn)里面存在大量的取模和坐標(biāo)映射操作（nn_bw_compute_source_index_fn）以及循環(huán)統(tǒng)計(jì)貢獻(xiàn)等。對(duì)于深度學(xué)習(xí)來說，UpsampleNearest最常用的其實(shí)就是2倍上采樣，比如Unet和YOLOv5，所以我們完全可以針對(duì)這種情況寫一個(gè)特化的Kernel，很輕量的來完成2倍上采樣的計(jì)算。下面展示OneFlow中針對(duì)2倍上采樣的優(yōu)化（代碼見：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/UpsampleNearest2D/upsample_nearest_2d.cu#L16-L63）

// CUDA: grid stride looping
#define CUDA_1D_KERNEL_LOOP(i, n)                                                                 \
  for (int32_t i = blockIdx.x * blockDim.x + threadIdx.x, step = blockDim.x * gridDim.x; i < (n); \
       i += step)

// Upsample Nearest2D Kernel is copyed from https://github.com/Oneflow-Inc/oneflow/blob/master/oneflow/user/kernels/upsample_nearest_kernel.cu#L78
template<typename T>
struct alignas(2 * sizeof(T)) Pack2X {
  T x;
  T y;
};

template<typename T>
__global__ void UpsampleNearest2D2XForward(const int32_t in_elem_cnt, const T* in_dptr,
                                           const int32_t in_height, const int32_t in_width,
                                           T* out_dptr) {
  const int32_t in_hw_size = in_width * in_height;
  CUDA_1D_KERNEL_LOOP(index, in_elem_cnt) {
    const T in_value = in_dptr[index];
    const int32_t nc_idx = index / in_hw_size;
    const int32_t hw_off = index - nc_idx * in_hw_size; // 這里是優(yōu)化掉昂貴的取模運(yùn)算
    const int32_t h = hw_off / in_width;
    const int32_t w = hw_off - h * in_width;
    Pack2X<T> out_value{in_value, in_value};
    Pack2X<T>* out_pack_dptr = reinterpret_cast<Pack2X<T>*>(out_dptr);
    out_pack_dptr[nc_idx * in_hw_size * 2 + h * 2 * in_width + w] = out_value;
    out_pack_dptr[nc_idx * in_hw_size * 2 + (h * 2 + 1) * in_width + w] = out_value;
  }
}

template<typename T>
__global__ void UpsampleNearest2D2XBackward(const int32_t in_elem_cnt, const T* dy_dptr,
                                            const int32_t dx_height, const int32_t dx_width,
                                            T* dx_dptr) {
  const int32_t dx_hw_size = dx_height * dx_width;
  CUDA_1D_KERNEL_LOOP(index, in_elem_cnt) {
    T dx_value = 0.0;
    const int32_t nc_idx = index / dx_hw_size;
    const int32_t dx_hw_off = index - nc_idx * dx_hw_size;
    const int32_t dx_h = dx_hw_off / dx_width;
    const int32_t dx_w = dx_hw_off - dx_h * dx_width;
    const Pack2X<T>* dy_pack_dptr = reinterpret_cast<const Pack2X<T>*>(dy_dptr);
    const Pack2X<T> dy_pack_value1 =
        dy_pack_dptr[nc_idx * dx_hw_size * 2 + dx_h * 2 * dx_width + dx_w];
    const Pack2X<T> dy_pack_value2 =
        dy_pack_dptr[nc_idx * dx_hw_size * 2 + (dx_h * 2 + 1) * dx_width + dx_w];
    dx_value += dy_pack_value1.x;
    dx_value += dy_pack_value1.y;
    dx_value += dy_pack_value2.x;
    dx_value += dy_pack_value2.y;
    dx_dptr[index] = dx_value;
  }
}

這個(gè)地方比較好理解，我們以前向的UpsampleNearest2D2XForward為例，當(dāng)我們對(duì)一個(gè)的矩陣進(jìn)行2倍上采樣時(shí)，可以獲得大小的輸出Tensor，那么輸入和輸出的對(duì)應(yīng)關(guān)系如下圖所示：

也就是輸入的(0, 0)位置對(duì)應(yīng)來輸出的(0, 0), (0, 1), (1, 0), (1, 1)的位置。也就是一個(gè)輸入的元素其實(shí)是對(duì)應(yīng)來輸出的4個(gè)元素，并且這4個(gè)元素一定是相鄰的2行或2列。所以我們可以使用Pack技術(shù)只用2次賦值就完成輸出Tensor對(duì)應(yīng)位置元素的填寫，進(jìn)一步提升全局內(nèi)存訪問的帶寬。

我這里直接使用 oneflow 的腳本對(duì)這兩個(gè) kernel 進(jìn)行進(jìn)行 profile ：

import oneflow as flow

x = flow.randn(16, 32, 80, 80, device="cuda", dtype=flow.float32).requires_grad_()

m = flow.nn.Upsample(scale_factor=2.0, mode="nearest")

y = m(x)
print(y.device)
y.sum().backward()

下面展示了在 A100 上調(diào)優(yōu)前后的帶寬占用和計(jì)算時(shí)間比較：

可以看到基于 oneflow upsample_nearest2d 的前后向的優(yōu)化 kernel 可以獲得更好的帶寬利用率和性能。

模板代碼和用法詳見：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/blob/master/UpsampleNearest2D/upsample_nearest_2d.cu

0x4. 總結(jié)

本篇文章推薦和講解一下OneFlow ElementWise模板，F(xiàn)astAtomicAdd，OneFlow UpsampleNearest2d模板的用法以及原理，并將其整理為最小的可以白嫖的頭文件。相關(guān)代碼請(qǐng)?jiān)L問 https://github.com/BBuf/how-to-optim-algorithm-in-cuda 這里獲得。