在CUDA中查找第一个非零元素

#include <cuda_runtime.h>
#include <stdio.h>

__global__ void findFirstNonZeroElement(const int* data, int size, int* result)
{
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    
    if (tid < size && *result == -1)
    {
        if (data[tid] != 0)
        {
            atomicCAS(result, -1, tid);
        }
    }
}

int main()
{
    const int size = 100;
    int data[size];
    int result = -1;
    
    // Initialize data
    for (int i = 0; i < size; i++)
    {
        data[i] = i % 10; // Set some elements to zero
    }
    
    int* d_data;
    int* d_result;
    
    // Allocate memory on GPU
    cudaMalloc((void**)&d_data, size * sizeof(int));
    cudaMalloc((void**)&d_result, sizeof(int));
    
    // Copy data from host to device
    cudaMemcpy(d_data, data, size * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_result, &result, sizeof(int), cudaMemcpyHostToDevice);
    
    // Launch kernel
    int blockSize = 256;
    int numBlocks = (size + blockSize - 1) / blockSize;
    findFirstNonZeroElement<<<numBlocks, blockSize>>>(d_data, size, d_result);
    
    // Copy result from device to host
    cudaMemcpy(&result, d_result, sizeof(int), cudaMemcpyDeviceToHost);
    
    // Free memory on GPU
    cudaFree(d_data);
    cudaFree(d_result);
    
    if (result != -1)
    {
        printf("The first non-zero element is at index %d\n", result);
    }
    else
    {
        printf("No non-zero element found\n");
    }
    
    return 0;
}