原文链接

https://developer.nvidia.com/blog/cuda-pro-tip-optimized-filtering-warp-aggregated-atomics/

场景

将 src 中大于 0 的数字移动到 dst 中去，并返回新数组的元素数量。和 C++ 的 std::copy_if 相似。

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 the number of elements copied
  return nres;
}

如果用全局内存实现：

__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];
}

每个参与的进程用 atomicAdd() 占位得到自己的坐标（atomicAdd() 返回原子操作之前的旧值），然后将结果写进去。缺点是全局内存上原子操作开销大，尤其是在竞争激烈的时候。

用共享内存实现略。在原文写了一块代码，代码不仅复杂，效果也仅仅比全局内存好了一点点（都很差）。

Warp-aggregated Atomics

主要流程如下，和全局内存上的实现很像，只是 atomicAdd() 改成了 atomicAggInc()：

__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];
}

重点是怎么把 atomicAggInc() 写的快一点：

1. __activemask() 获取 warp 内运行到这里的线程数量（每个线程对应一个 lane）。
2. 然后选出 leader。
3. 计算 rank，即当前线程在集合中的排行。__lanemask_lt() 返回的是一个掩码，和 active 做与运算可以得到比当前 lane 还小的、active 中包含的 lane 的集合。再求 1 的数量就得到了 rank。
4. Leader 的 rank 为 0，我们只让 leader 去做原子加法，在全局的数组中给当前的集合占一小块位置。然后有了起始位置和 rank，每个线程就知道自己的位置对应哪里了。
5. 在返回之前，还要用 __shfl_sync() 将结果广播到整个 warp，否则在其他线程看来，warp_res 是没有初始化的。

__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;
}

选 leader 的细节：

1. __ffs() 找到元素中最低位 1 的位置，参数为 0 则返回 0，否则最低位的 1 从 1 开始计数。因为当前运行到这里的线程至少有 1 个（线程自己），所以 __ffs() 返回值必然大于 0。
2. 接下来计算要写几个元素进去（dst），这里的逻辑是用 __popc() 对 active 的 1 计数。所以调用者一定要保证该线程确实要写东西进来，才调用 atomicAggInc()。这一点是在前面的代码中用 if 判断保证的。

为什么不借助共享内存完成数据交换？Warp 的原语能完成寄存器之间的数据交换，不用先写后读，也不需要经过共享缓存，因而效率更高。

atomicAggInc() 里面也用到了 atomicAdd()，为什么前者更快？前者一次将一个 warp 的数据增量计算出来加到了全局计数变量上，使得竞争数减少了若干倍（最好的情况下竞争可以减少为 1/32）。

其他

在 NVIDIA 的博客 https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/ 中，atomicAggInc() 的实现多了一行：

int mask = __match_any_sync(__activemask(), (unsigned long long)ptr);

ptr 是全局计数变量的地址，atomicAggInc() 用 __match_any_sync() 对当前运行到这里的集合再次分区，找到当前 ptr 值一样的那些线程。在上文的例子中，每个 warp 中的计数变量的地址是相同的，所以没有这一步操作。