GPU葵花宝典:当原子操作遇上warp：GPU界的"拼团"大法

不会弹吉他的布鲁克
138 阅读4分钟

现在NVCC编译器已能自动实现线warp级原子操作聚合，躺着也能享受性能飙升！不过想了解原理的小伙伴们，快搬小板凳来听书，这招在其他场景也能大显身手哦

问题场景：数组元素大筛选
想象你正在处理一个超长数组，要从中挑出所有正数放进新数组。CPU版的常规操作就像超市收银员逐个结账：

int filter(int *dst, const int *src, int n) {
  int nres = 0;
  for (int i = 0; i < n; i++)
    if (src[i] > 0)
      dst[nres++] = src[i];  // 收银员默默记账
  return nres;
}

但到了GPU的并行世界，事情就刺激了——所有线程都抢着要当"记账员"！

初代GPU方案：全员抢柜台

__global__ void filter_k(int *dst, int *nres, const int *src, int n) {
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  if(i < n && src[i] > 0)
    dst[atomicAdd(nres, 1)] = src[i]; // 万人血书抢柜台
}

结果可想而知：当半数元素需要记账时，性能直接"跳水",GPU 并行世界中,原子操作就像春运抢票，人越多越挤爆服务器！

image.png

优化方案 1: 执行的原子操作数量或数组中正元素的比例与带宽成反比,网格中所有读取正元素的线程都会增加一个单一的计数器 nres 。根据正元素的数量，这可能是大量的线程。因此， atomicAdd() 的碰撞程度很高，这限制了性能,就像春运2亿人在抢唯一的一张票,因此在不能增加票的前提下,减少买票行为的请求就是最优的解决办法了,如果客官你着急回去办事,请跳到最后看终极方案,如果你还有雅兴,先喝口茶润润嗓子

下面我们来利用共享内存和原子操作在 GPU 上实现一种分块并行过滤的方法，有效避免了多线程同时写入全局内存时可能出现的竞争条件，同时充分利用了 CUDA 的并行计算能力

__global__ 
void filter_shared_k(int *dst, int *nres, const int* src, int n) {
  __shared__ int l_n;
  int i = blockIdx.x * (NPER_THREAD * BS) + threadIdx.x;//每个线程处理 n 个元素

  for (int iter = 0; iter < NPER_THREAD; iter++) {
    // zero the counter
    if (threadIdx.x == 0)
      l_n = 0;
    __syncthreads();

    // get the value, evaluate the predicate, and
    // increment the counter if needed
    int d, pos;

    if(i < n) {
      d = src[i];
      if(d > 0)
        pos = atomicAdd(&l_n, 1);
    }
    __syncthreads();

    // leader increments the global counter
    if(threadIdx.x == 0)
      l_n = atomicAdd(nres, l_n);
    __syncthreads();

    // threads with true predicates write their elements
    if(i < n && d > 0) {
      pos += l_n; // increment local pos by global counter
      dst[pos] = d;
    }
    __syncthreads();

    i += BS;
  }
}

本方案采用基于共享内存的原子操作, 基于分块的思想利用一个线程连续处理多个元素,且对共享内存执行 atomic 原子操作进行性能提升

另一种方法是首先使用并行前缀和来计算每个元素的输出索引。Thrust 的 copy_if() 函数使用了这种方法的优化版本。图 2 展示了这两种方法在 Kepler K80 上的性能。尽管共享内存原子操作可以提高过滤性能，但仍然只比原始方法快 1.5 倍左右。原子操作仍然是瓶颈，因为操作次数没有变化。对于高过滤比例，Thrust 优于两种方法，但会产生较大的前期成本，而这种成本对于小过滤比例来说是无法摊薄的

image.png

性能救星： warp 内组团买单
核心思想很简单——让32个线程先内部"拼单"，再派代表统一结账：

  1. 选个团长（线程束中最低ID的活跃线程）
  2. 统计本团购买总数
  3. 团长单枪匹马执行原子操作
  4. 全员按团购量分配位置

代码演绎

// 使用Cooperative Groups的优雅版
__device__ int atomicAggInc(int *ctr) {
  auto g = coalesced_threads();  // 召唤拼单小分队
  int warp_res;
  if(g.thread_rank() == 0)
    warp_res = atomicAdd(ctr, g.size());  // 团长结账
  return g.shfl(warp_res, 0) + g.thread_rank(); // 分发战利品
}

// 硬核玩家的原始操作版
__device__ int atomicAggInc(int *ctr) {
  unsigned int active = __activemask();
  int leader = __ffs(active) - 1;  // 选团长：最靓的仔
  int change = __popc(active);     // 拼单总数
  unsigned int rank = __popc(active & __lanemask_lt());
  int warp_res;
  if(rank == 0)
    warp_res = atomicAdd(ctr, change);  // 团长出击
  warp_res = __shfl_sync(active, warp_res, leader);
  return warp_res + rank;  // 人手一个VIP号
}

__global__ void filter_k(int *dst, const int *src, int n) {
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  if(i >= n) 
    return; //线程非活跃
  if(src[i] > 0)
    dst[atomicAggInc(nres)] = src[i];
}

其他基于协作组的提速操作可以参考官方仓库的案例 GPUCache

性能实测：质的飞跃 这波操作直接让原子操作量减少32倍！,性能提升21倍

image.png

image.png

更妙的是，CUDA 9之后开始编译器已自动实现这波操作,但是手动实现有助于我们了解其内在机制和原理

  • 学习原子操作优化原理
  • 适用于其他需要自定义原子逻辑的场景
  • 在老版不能升级的CUDA环境（如CUDA 8）中手动优化

各代GPU性能对比
从Kepler到Volta架构（图4-6），自动优化的原子操作性能已媲美手动优化版。不过技术控们要记住：

  • Pascal架构开始硬件支持原子操作优化
  • Volta架构的独立线程调度需要特别注意活跃掩码
  • 共享内存原子操作在新时代GPU中反而可能拖后腿
avatar
文章
阅读
粉丝