乐趣区

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

撰文|郑泽康

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 中,整数除法的计算代价是比拟低廉的(https://docs.nvidia.com/cuda/…)对于计算指令耗时这一章中有提到:

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 最新版本:
https://github.com/Oneflow-In…

退出移动版