读英伟达博客:CUDA Wrap 级原语

原文

https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/

笔记

Warp 内规约

#define FULL_MASK 0xffffffff
for (int offset = 16; offset > 0; offset /= 2)
    val += __shfl_down_sync(FULL_MASK, val, offset);

先在 warp 内用原语规约,比在 block 级别用共享内存规约快很多。(虽然也有 __reduce_add_sync() 函数,但是截至 2024 年 9 月 5 日只支持 unsigned 和 int 类型,给浮点数做规约就要用 __shfl_xx_sync() 系列。)

CUDA 9 的相关 API

直接从原博客摘录:

  1. Synchronized data exchange: exchange data between threads in warp. 形式为 __xx_sync(mask, val)
    • __all_sync, __any_sync, __uni_sync, __ballot_sync
    • __shfl_sync, __shfl_up_sync, __shfl_down_sync, __shfl_xor_sync
    • __match_any_sync, __match_all_sync
  2. Active mask query: returns a 32-bit mask indicating which threads in a warp are active with the current executing thread.
    • __activemask
  3. Thread synchronization: synchronize threads in a warp and provide a memory fence.
    • __syncwarp

说明:

  1. __syncwarp() 只保证在运行此函数的时候在 warp 内达成同步,但不保证同步后的下一条语句依然由 warp 内所有的线程参与执行。应该只能用来同步到共享内存的写入和读取?
  2. __xx_sync(mask, val) 不等于 __xx(val)__syncwarp() 的组合。
  3. __balloc_sync(mask, pred) 是从 mask 中找出 pred 非零的、还未退出的线程的集合,返回方式为 32 位掩码。
  4. __match_any_sync(mask, val) 是把 val 相同的分区到一起,返回当前区的 mask。
  5. __match_all_sync(mask, val, pred) 要么返回 mask(mask 里的所有线程都分区到一起),要么返回 0。传出参数 pred 根据是否成功设置成 true 或者 false。

Warp 内同步

void __syncwarp(unsigned mask=FULL_MASK);

考虑不使用 __shfl_down_sync(),而是用类似共享内存内累加的方式完成 warp 内的累加。

错误的 warp 内累加

unsigned tid = threadIdx.x;

// Incorrect use of __syncwarp()
shmem[tid] += shmem[tid+16]; __syncwarp();
shmem[tid] += shmem[tid+8];  __syncwarp();
shmem[tid] += shmem[tid+4];  __syncwarp();
shmem[tid] += shmem[tid+2];  __syncwarp();
shmem[tid] += shmem[tid+1];  __syncwarp();

✔ 正确的 warp 内累加:

unsigned tid = threadIdx.x;
int v = 0;

v += shmem[tid+16]; __syncwarp();
shmem[tid] = v;     __syncwarp();
v += shmem[tid+8];  __syncwarp();
shmem[tid] = v;     __syncwarp();
v += shmem[tid+4];  __syncwarp();
shmem[tid] = v;     __syncwarp();
v += shmem[tid+2];  __syncwarp();
shmem[tid] = v;     __syncwarp();
v += shmem[tid+1];  __syncwarp();
shmem[tid] = v;

2025/3/25 我试了一下用 if (tid < stride) shmem[tid] += shmem[tid+stride]; __syncwarp(); 也是能工作的。这个和 blockreduce 的模式一样,用 if 条件可以保证正在被读的对象不会被写。

机会主义的 warp 级编程

原文给了一个把 per-thread 原子加改写成 per-warp 原子加的代码:

// increment the value at ptr by 1 and return the old value
__device__ int atomicAggInc(int *ptr) {
    // group threads by ptr value
    int mask = __match_any_sync(__activemask(), (unsigned long long)ptr);
    int leader = __ffs(mask) – 1;    // select a leader
    int res;
    if(lane_id() == leader)                  // leader does the update
        res = atomicAdd(ptr, __popc(mask));
    res = __shfl_sync(mask, res, leader);    // get leader’s old value
    // let every thread THINK that they did the atomic op separately
    return res + __popc(mask & ((1 << lane_id()) – 1)); //compute old value
}

这段代码写的比 读英伟达博客:Warp 聚合原子操作 —— 以过滤为例 考虑得还要周到一点。代码解释可以参考刚刚给出的另外一篇博客。

其他建议

文章后面的内容大概意思是建议全部换用新的 __xx_sync() API,不要再用旧的。文章也讲了一下原来的写法中存在的风险。主要原因还是原来的 __xx() API 不接受掩码参数,不会(对指定的线程)显式同步,因此依赖于运行时逐语句的同步执行。但是新的 GPU 上的 warp 调度优化越来越多,按同一个步调执行每条语句这个条件不一定还能满足。