关于cuda:CUDA优化之PReLU性能调优

撰文|郑泽康

InsightFace 模型里大量应用了 PReLU 激活函数，而 PReLU 的工作模式有两种：

PReLU(1)，此时权重 alpha 的形态为(1,)，等价于一个 Elementwise 操作。
PReLU(channels)，此时权重 alpha 的形态为 (channels,)，和输出特色(N, C, H, W) 中 C 的大小是对应的。此时 PReLU 等价于一个 Binary Broadcast 操作。

InsightFace 模型里的 PReLU 工作模式是第二种，之前曾经介绍过 CUDA Elementwise 操作优化，而在 Broadcast 情景下也存在肯定的优化机会。

一个奢侈实现的思维就是在循环外部，依据以后元素的索引，推算出该元素对应须要应用的 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 中，整数除法的计算代价是比拟低廉的（https://docs.nvidia.com/cuda/…）对于计算指令耗时这一章中有提到：

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

整数除法，取余操作会被编译成多达 20 条指令。而咱们这里计算 alpha 的索引的时候，别离用到一次除法，一次取余，占整个 Kernel 的次要计算量，上面咱们将用向量化的思路来进步读写带宽的同时，缩小整数除法，取余的计算次数。

咱们思考一个比较简单的例子，输出为(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 整除时，只能退回到奢侈实现上。

在 A100-40GB 测试机器上，咱们对 Insightface 波及到的 Tensor 形态，与 PyTorch 实现进行比拟，测试数据如下：

通过优化 PReLU 的 OneFlow，在大部分状况下均有比 PyTorch 靠近 2 倍的当先劣势，在最初一种状况因为形态较为非凡，无奈利用向量化的优化，所以体现与 PyTorch 持平。

欢送下载体验 OneFlow v0.7.0 最新版本：
https://github.com/Oneflow-In…

关于cuda:CUDA优化之PReLU性能调优

奢侈实现

Pack 向量化优化

基准测试