CUDA：在 O(1) 时间内使用扭曲中的所有线程设置第 N 个位索引

Question

我有一个 32 位位掩码，包含一组有效项目。
我想从该位掩码中提取有效条目的索引作为列表。 假设我使用选票获得了位掩码，并且我想知道有效线程的 id 作为顺序列表，例如：

const auto bitmask = __ballot_sync(-1u, isValid);
if (threadIdx.x == 0) { printf("Bitmask = $%x\n); } 

例如，如果掩码是

0b100010001

我想提取三项 0、4、8。
如何使用扭曲中的所有线程在恒定时间内（无循环）提取 int32 中设置位的位索引？

Answer 1

CudaaduC 论坛帖子的简化版本只需几条说明即可工作：

auto lanemask_lt() {
  uint32_t result;
  asm ("mov.u32 %0, %%lanemask_lt;" : "=r"(result));
  //or: -1u >> (32 - (threadIdx.x % 32));
  return result
}

auto lanemask_eq() {
  uint32_t result;
  asm ("mov.u32 %0, %%lanemask_eq;" : "=r"(result));
  //or: return 1 << (threadIdx.x % 32);
  return result
}

通道掩码是读取特殊 GPU 寄存器的内在指令。 这些与 popcount 相结合来获取位置。

void PrintValidIndices() {
  assert(threadIdx.x < 32);
  const auto isValid = ((threadIdx.x % 3) == 0);
  __shared__ int ValidIndex[32];
  ValidIndex[threadIdx.x] = -1;
  const auto pos = __popc(Mask & lanemask_lt());
  //use if you don't have an isValid flag from before
  //const auto isValid = (Mask & lanemask_eq()); 
  __syncwarp();
  if (isValid) { ValidIndex[pos] = threadIdx.x; }
  __syncwarp();   
  printf("ValidIndex[%i] = %i\n", threadIdx.x, ValidIndex[threadIdx.x]);
}