如何在 CUDA 中的条件 while 语句中同步 warp 中的线程？

Question

假设我们有以下代码：

while (condition) {
  ...

  for (uint32_t gap = x >> 1; gap > 0; gap >>= 1) {
    val += __shfl_down_sync(mask, val, gap);
  }

  if (warpLane == 0)
    atomicAdd(&global_memory[threadIdx.x], val);

  ...
}

在这种情况下，如果warp中的线程按以下顺序进入while循环：

全部 32 个线程，全部 32 个线程，只有 16 个线程。

如何获得参与 while 循环语句的线程掩码？

根据https://devblogs.nvidia.com/using-cuda-warp-level-primitives 中描述的指南，以下代码可能会导致未定义的行为：

while (condition) {
  uint32_t active = __activemask();
  for (uint32_t gap = x >> 1; gap > 0; gap >>= 1) {
    val += __shfl_down_sync(active, val, gap);
  }

  if (warpLane == 0)
    atomicAdd(&global_memory[threadIdx.x], val);

  ...
}

根据指南，__activemask() 可能不会像我预期的那样生成掩码。

根据上述指南，以下也会导致未定义的行为：

while (condition) {
  uint32_t active = __activemask();
  for (uint32_t gap = x >> 1; gap > 0; gap >>= 1) {
    val += __shfl_down_sync(active, val, gap);
  }

  if (warpLane == 0)
    atomicAdd(&global_memory[threadIdx.x], val);

  ...
  __warpsync(active);
}

那么，如何正确获取口罩呢？

Answer 1

您可以使用cooperative groups 喜欢：

#include <cooperative_groups.h>
namespace cg = cooperative_groups;

while (condition) { 
...
auto active = cg::coalesced_threads(); // this line can be moved out of while if the condition does not cause thread divergence

 for (uint32_t gap = x >> 1; gap > 0; gap >>= 1) { 
        //val += __shfl_down_sync(mask, val, gap);
        val += active.shfl_down(val, gap);
 }
 if (warpLane == 0)
    atomicAdd(&global_memory[threadIdx.x], val); 

... 
}

如果你想自己生成面具并做老式的你可以使用：

uint32_t FullMask = 0xFFFFFFFF;
uint32_t mask =  __ballot_sync(FullMask, someCondition);

但是，如果您在代码中有进一步的分支，则必须始终在分支之前跟踪 mask，并在 ballot 中使用它而不是 FullMask。所以分支之前的第二次更新将是：

uint32_t newMask =  __ballot_sync(mask, someNewCondition);