CUDA中如何实现布尔索引数组到整数索引数组的转换

核心逻辑说明 #

编写该转换逻辑的核心卡点是：每个命中真值的位置，需要明确它在所有真值序列中的排序位置，也就是输出数组的下标——这个步骤的本质是布尔掩码的并行前缀求和问题，完全可以全流程并行实现，不需要串行遍历。

自定义Kernel实现方案 #

这里提供两种可直接落地的实现思路，适配不同的开发和性能需求：

方案1：原子计数器极简实现（适合快速开发） #

不需要复杂的同步或扫描逻辑，直接在原有代码框架上改造即可：

1. 提前在设备内存上分配一个整型计数器，初始值为0，专门记录已经写入输出数组的有效元素个数
2. 每个线程检测到负责位置的布尔值为真时，通过原子操作获取当前计数器的旧值作为写入下标，同时把计数器的值加1

改造后的Kernel代码：

// write_offset 为设备内存上的计数器，需提前用cudaMemset设置为0
__global__ void transfer(int ni1, const bool* idx1, int* idx2, int* write_offset)
{
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  if (i < ni1 && idx1[i])
  {
    // 原子返回操作前的计数器值，避免多线程写冲突
    int write_pos = atomicAdd(write_offset, 1);
    idx2[write_pos] = i; // 如果需要从1开始计数的索引，替换为 i+1
  }
}

• 优势：代码改动量极小，无额外同步开销，适配绝大多数中等规模数据场景
• 劣势：高真值占比、超大规模数组场景下，原子操作的串行开销会拉低性能

方案2：并行前缀和实现（适合高性能需求） #

完全消除原子操作，做到无冲突全并行执行：

1. 对原始布尔数组做转换：真值记为1，假值记为0
2. 计算转换后数组的包含式前缀和：也就是每个位置i对应的值，代表从数组开头到i位置一共累计了多少个真值
3. 每个线程检测到负责位置的布尔值为真时，直接把当前位置i写入到前缀和值 - 1对应的输出位置即可

对应样例数组的计算过程：

输入布尔数组：[0,1,0,0,1,1,0,1]
转换为数值数组：[0,1,0,0,1,1,0,1]
计算后的前缀和数组：[0,1,1,1,2,3,3,4]
最终输出的0基整数索引：[1,4,5,7]，和转换逻辑完全匹配

CUDA 11及以上版本可以直接调用设备端标准库的扫描接口实现前缀和，不需要手动编写分块扫描逻辑：

#include <cuda/std/scan>
#include <cuda/std/execution>
// 步骤1：把布尔数组转成可计算的数值数组
cuda::std::transform(
  cuda::std::device, idxA, idxA + ni1, temp_num_array,
  []__device__(bool mask_val){ return mask_val ? 1 : 0; }
);
// 步骤2：计算前缀和
cuda::std::inclusive_scan(
  cuda::std::device, temp_num_array, temp_num_array + ni1, prefix_sum_array
);
// 步骤3：遍历写入结果
__global__ void write_result(int ni1, const bool* idx1, const int* prefix_sum, int* idx2)
{
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  if (i < ni1 && idx1[i])
  {
    idx2[prefix_sum[i] - 1] = i;
  }
}

Thrust库实现方案 #

Thrust提供了现成的copy_if接口，不需要手动编写Kernel即可完成转换，性能经过充分优化：

1. 先生成和输入布尔数组等长的连续整数序列，序列值为每个位置的索引
2. 以布尔数组为判断条件，把序列中满足“布尔值为真”的元素拷贝到输出数组

#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <thrust/count.h>
#include <thrust/sequence.h>

// 输入布尔数组
thrust::device_vector<bool> d_bool_mask = {0,1,0,0,1,1,0,1};
// 生成连续索引序列，默认从0开始计数
thrust::device_vector<int> d_index_seq(d_bool_mask.size());
thrust::sequence(d_index_seq.begin(), d_index_seq.end(), 0);
// 提前根据真值个数分配输出数组空间
int valid_count = thrust::count(d_bool_mask.begin(), d_bool_mask.end(), true);
thrust::device_vector<int> d_int_idx(valid_count);
// 执行条件拷贝
thrust::copy_if(
  d_index_seq.begin(), d_index_seq.end(),
  d_bool_mask.begin(),
  d_int_idx.begin(),
  []__device__(bool mask_val){ return mask_val; }
);

如果需要得到示例中从1开始计数的输出结果，只需要在调用thrust::sequence时把起始偏移设置为1即可。

内容的提问来源于stack exchange，提问作者Wei Zhang