如何使用CUDA结合原子操作从设备数组生成值为1的元素索引数组？

核心逻辑 #

在并行设备（如GPU）上，让每个线程负责处理原数组的单个元素：

• 若当前元素值为1，通过原子递增操作获取结果数组的当前写入位置，再将该元素的索引存入对应位置。
• 原子操作能避免多线程竞争计数器，保证索引写入不会出现混乱或覆盖。

具体代码实现 #

1. 核函数（设备端执行） #

__global__ void collectIndices(const int* input, int inputSize, int* output, int* counter) {
    // 计算当前线程处理的原数组索引
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 检查索引合法性，且当前元素为1时执行写入
    if (idx < inputSize && input[idx] == 1) {
        // 原子递增计数器，返回递增前的旧值作为写入位置
        int writePos = atomicAdd(counter, 1);
        // 将当前索引存入结果数组对应位置
        output[writePos] = idx;
    }
}

2. 主机端调用示例 #

#include <iostream>
#include <vector>
#include <cuda_runtime.h>

int main() {
    // 主机端原数组
    std::vector<int> hostInput = {0,1,0,0,1,1,0};
    int inputLen = hostInput.size();
    
    // 结果数组预分配最大长度（最坏情况原数组全为1）
    std::vector<int> hostOutput(inputLen);
    int hostCounter = 0;

    // 设备端内存分配
    int* devInput;
    int* devOutput;
    int* devCounter;
    cudaMalloc(&devInput, inputLen * sizeof(int));
    cudaMalloc(&devOutput, inputLen * sizeof(int));
    cudaMalloc(&devCounter, sizeof(int));

    // 数据从主机拷贝到设备
    cudaMemcpy(devInput, hostInput.data(), inputLen * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(devCounter, &hostCounter, sizeof(int), cudaMemcpyHostToDevice);

    // 启动核函数（按256线程/块的标准配置）
    int blockSize = 256;
    int gridSize = (inputLen + blockSize - 1) / blockSize;
    collectIndices<<<gridSize, blockSize>>>(devInput, inputLen, devOutput, devCounter);
    cudaDeviceSynchronize();

    // 将结果和计数器值拷贝回主机
    cudaMemcpy(&hostCounter, devCounter, sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(hostOutput.data(), devOutput, hostCounter * sizeof(int), cudaMemcpyDeviceToHost);

    // 输出结果
    std::cout << "收集到的索引：";
    for (int i = 0; i < hostCounter; ++i) {
        std::cout << hostOutput[i] << " ";
    }
    std::cout << std::endl;

    // 释放设备内存
    cudaFree(devInput);
    cudaFree(devOutput);
    cudaFree(devCounter);

    return 0;
}

关键细节说明 #

• 原子操作atomicAdd：这是核心，它会原子性地完成「读取计数器值→加1→写回」的操作，返回的旧值就是当前线程的合法写入位置，彻底避免多线程冲突。
• 结果数组长度：因为无法提前预知原数组中1的数量，所以结果数组需要分配与原数组等长的空间，最终通过counter的值确定有效元素的实际个数。
• 其他并行设备适配：如果是OpenCL等其他并行框架，思路完全一致，只是原子操作的函数名可能不同（比如OpenCL用atomic_inc），核心都是通过原子计数安全分配写入位置。

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