目录CUDA 线程同步协作组同步grid_group::sync() 源码分析实现自定义跨块同步测试小结笔者的毕设项目与 CUDA 相关最近需要实现设备端跨线程块同步操作。查阅了相关 API 后发现有一个集群Cluster同步操作可用但是集群需要计算能力 9.0而服务器配备 RTX 4090计算能力 8.9一些云算力提供 V100计算能力 7.0、T4计算能力 7.5均不达要求。最终参考了协作组的网格同步函数实现了兼容低计算能力的自定义跨线程块同步。CUDA 线程同步传统 CUDA 编程模型是面向线程的单指令多线程Single-Instruction Multi-Thread, SIMT模型线程结构从高到低分为网格grid、线程块block、线程束warp、线程thread。启动一个核函数kernel时会把它分配到一个流多处理器Streaming Multiprocessor, SM上的多个流处理器Streaming Processor, SP上。受 SM 的寄存器等资源限制GPU 会以 32 个线程组成一个线程束warp进行调度。线程束是调度和运行的基本单位。Volta 架构计算能力 7.0之前线程束内的所有线程共享程序计数器PC并行执行相同的指令从 Volta 架构开始每个束内线程拥有独立的 PC实现了更灵活的 SIMT 模型。当一个并行任务存在多个步骤及数据依赖时需要等待前一步所有线程执行完毕后才能执行后一步任务这需要线程同步。CUDA 支持三种级别的线程同步网格同步当一个网格的所有线程均完成一个步骤后进行同步。主机端可以将任务按照步骤拆分为多个 kernel 依次执行实现网格同步设备端则需要借助协作线程组执行grid.sync()实现。线程块同步当一个线程块内的所有线程均完成一个步骤后进行同步。设备端通过__syncthreads()原语或者__barrier_sync()实现。线程束同步当一个线程束内的所有线程均完成一个步骤后进行同步。Volta 架构之前线程束始终同步执行之后需要通过设备端原语__sync_warp()实现。当并行处理批量数据时如果一条数据分配到一个线程块那么内部使用__syncthreads()即可。然而面对小批量数据时这种分配方式导致 SM 利用率低下通常需要将单个数据分配到多个线程块以充分利用 SM 资源。协作组同步CUDA 9 引入协作组Cooperative Groups用于组织通信线程组。协作组提供了更多层级的通信线程组划分与同步方法。协作组相关功能需要通过cooperative_groups.h引入除了提供了网格、线程块、线程束级别的协作组同步功能还支持集群、线程块分片thread block tile等介于不同级别之间的协作组。使用较多的是网格同步通过cooperative_groups::this_grid()取得当前网格协作组后调用grid_group::sync()方法可进行网格同步。集群介于网格和线程块之间由多个线程块组成通过cooperative_groups::this_cluster().sync()进行集群同步。CUDA API 中网格同步使用软件实现而集群同步使用__cluster_barrier_arrive()和__cluster_barrier_wait()等原语实现要求计算能力 9.0 及以上。为了实现低计算能力的类似集群同步功能我们可以参考网格同步来软件实现。grid_group::sync() 源码分析以 CUDA 12.8 版本 API 源码为例查看grid_group::sync()源码_CG_QUALIFIER void sync() const {if (!is_valid()) {_CG_ABORT();}details::grid::sync(_data.grid.gridWs-barrier);}is_valid()用于验证_data.grid.gridWs是否非空_data.grid.gridWs-barrier是当前网格的一个屏障变量。继续查看details::grid::sync()源码_CG_STATIC_QUALIFIER void sync(unsigned int *bar) {unsigned int token details::sync_grids_arrive(bar);details::sync_grids_wait(token, bar);}可以看到同步操作分为同步到达和同步两步。先查看details::sync_grids_arrive()源码typedef unsigned int barrier_t;_CG_STATIC_QUALIFIER bool is_cta_master() {return (threadIdx.x threadIdx.y threadIdx.z 0);}_CG_STATIC_QUALIFIER unsigned int sync_grids_arrive(volatile barrier_t *arrived) {unsigned int oldArrive 0;__barrier_sync(0); // 块内同步if (is_cta_master()) { // CTA内主线程unsigned int expected gridDim.x * gridDim.y * gridDim.z; // 待同步的线程数量。每个线程块取一个主线程即线程块的总数bool gpu_master (blockIdx.x blockIdx.y blockIdx.z 0); // 判断网格内主线程块unsigned int nb 1; // 屏障自增量if (gpu_master) {// 主块的自增量需要单独设置保证所有nb之和为0x80000000nb 0x80000000 - (expected - 1);}NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,// Barrier update with release; polling with acquire// SM 7.0 以上使用带有release语义的原子自加asm volatile(atom.add.release.gpu.u32 %0,[%1],%2; : r(oldArrive) : _CG_ASM_PTR_CONSTRAINT((unsigned int*)arrived), r(nb) : memory);,// Fence; barrier update; volatile polling; fence// 否则使用内存屏障原子相加__threadfence();oldArrive atomicAdd((unsigned int*)arrived, nb););}// 返回为自增前的屏障值return oldArrive;}在此我们可以看到该同步屏障变量的使用方式。sync_grids_arrive()通过__barrier_sync(0)进行线程块同步并通过is_cta_master()选择块内第一个线程来更新arrived屏障避免块内冲突对于GPU/网格内的主线程设置nb 0x80000000 - (expected - 1)以保证所有屏障自增量nb之和为 0x80000000。这样一来当屏障初始化为 0 时执行依次网格同步后取值为 0x80000000再次执行则恢复到 0屏障可以循环使用。_CG_STATIC_QUALIFIER bool bar_has_flipped(unsigned int old_arrive, unsigned int current_arrive) {return (((old_arrive ^ current_arrive) 0x80000000) ! 0);}_CG_STATIC_QUALIFIER void sync_grids_wait(unsigned int oldArrive, volatile barrier_t *arrived) {if (is_cta_master()) { // 仅 CTA 内主线程处理NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,// 计算能力不低于 7.0 时使用带有 acquire 语义得加载指令读取屏障变量并与屏障旧值比较符号unsigned int current_arrive;do {asm volatile(ld.acquire.gpu.u32 %0,[%1]; : r(current_arrive) : _CG_ASM_PTR_CONSTRAINT((unsigned int *)arrived) : memory);} while (!bar_has_flipped(oldArrive, current_arrive));,// 否则直接读取比较并用内存屏障保证内存可见性while (!bar_has_flipped(oldArrive, *arrived));__threadfence(););}__barrier_sync(0); // 线程块同步}sync_grids_wait()通过循环等待屏障变量最高位变动0→0x800000000x80000000→0实现同步。实现自定义跨块同步笔者遇到的场景是需要对blockIdx.y相同的多个blockIdx.x不同线程块进行同步。基于以上分析我们可以为blockIdx.y相同的每一组线程设定一个屏障变量 barrier然后将组内同步数量expected设为gridDim.x * gridDim.z、主块判断gpu_master设定为blockIdx.x 0 blockIdx.z 0即可达到效果。同步函数如下// sync_ctas.cuh#pragma once#define bar_has_flipped(a,b) ((((a)^(b))0x80000000)!0)__device__ __inline__ void sync_ctas(unsigned *bar) {unsigned nb blockIdx.x 0 ? 0x80000000 - (gridDim.x - 1) : 1;if (threadIdx.x 0) {unsigned oldarr, cuarr;#if __CUDA__ARCH__ 700asm __volatile__(atom.add.release.gpu.u32 %0,[%1],%2;:r(oldarr):l(bar),r(nb):memory);#else__threadfence();oldarr atomicAdd(bar, nb);#endifdo {#if __CUDA_ARCH__ 700asm __volatile__(ld.acquire.gpu.u32 %0,[%1];:r(cuarr):l(bar):memory);#elsecuarr *(volatile unsigned*)bar;#endif} while (!bar_has_flipped(oldarr, cuarr));#if __CUDA_ARCH__ 700__threadfence();#endif}__barrier_sync(0);}测试通过以下代码验证该同步功能。该代码实现了对一个数组的簇内分段重置并与网格同步进行对比。#include cooperative_groups.h#include sync_ctas.cuh#define CHECK(code) \{auto c(code);\if (c!cudaSuccess) {\fprintf(stderr, __FILE__:%d: %s\n, __LINE__, cudaGetErrorString(c));\abort();\}}namespace cg cooperative_groups;template bool grid_sync__device__ __inline__ void sync(unsigned *barrier) {if constexpr (grid_sync) {cg::this_grid().sync();} else {sync_ctas(barrier);}}template bool grid_sync, bool verifyfalse__global__ void reset(int *arr, unsigned l, unsigned *barriers) {auto cid blockIdx.x;auto id blockIdx.y;auto barrier barriers id;arr l * id;int val 0;// 重置10次for (int j0; j10; j) {// 多块协同重置arrfor (unsigned icid * blockDim.x threadIdx.x; i l; i blockDim.x * gridDim.x) {arr[i] val;}syncgrid_sync(barrier);if constexpr (verify) { // 单个线程校验同步逻辑正确if (threadIdx.x 0 cid 0 id 0) {for (unsigned i0; i l; i) {if (arr[i] ! val) {printf(err %d\n, id);break;}}}syncgrid_sync(barrier);}val ~val;}}int main() {int l 512;int *arr;unsigned *barr;cudaMalloc(arr, l*32*sizeof(int));cudaMalloc(barr, 32*sizeof(unsigned));cudaMemsetAsync(barr, 0, 32*sizeof(unsigned));cudaEvent_t st, mid, s2, end;dim3 gd(4,32,1), bd(128,1,1);cudaEventCreate(st);cudaEventCreate(mid);cudaEventCreate(s2);cudaEventCreate(end);// 校验效果void *args[] {arr, l, barr};cudaLaunchKernel(resetfalse, true, gd, bd, args, 0, 0);// 比较耗时cudaEventRecord(st);cudaLaunchKernel(resetfalse, gd, bd, args, 0, 0);cudaEventRecord(mid);cudaEventRecord(s2);cudaLaunchCooperativeKernel(resettrue, gd, bd, args, 0, 0);cudaEventRecord(end);cudaDeviceSynchronize();float elp;cudaEventElapsedTime(elp, st, mid);printf(ctas_sync %g ms\n, elp);cudaEventElapsedTime(elp, s2, end);printf(grid_sync %g ms\n, elp);cudaFree(arr);cudaFree(barr);cudaEventDestroy(st);cudaEventDestroy(mid);cudaEventDestroy(s2);cudaEventDestroy(end);CHECK(cudaGetLastError());return 0;}在 RTX 409024G显卡上测试多次运行结果如下rootnode12:~# ./a.outctas_sync 0.013312 msgrid_sync 0.047104 msrootnode12:~# ./a.outctas_sync 0.011264 msgrid_sync 0.012288 msrootnode12:~# ./a.outctas_sync 0.012288 msgrid_sync 0.017408 msrootnode12:~# ./a.outctas_sync 0.011264 msgrid_sync 0.012288 msrootnode12:~# ./a.outctas_sync 0.022528 msgrid_sync 0.070656 ms可以看到验证核函数没有输出错误字段说明同步功能正常输出中软件多块同步始终比网格同步更快约1.1x~3.5x证明了软件簇同步比网格同步更高效。