CUDA归约内核计算结果偏差2的问题排查求助

问题现象 #

基于归约实现的CUDA内核，在输入规模为2的幂时多数正常，但部分非2的幂规模（如N=(1<<13)+4）下计算结果偏差2。通过打印中间部分和发现：

threadIdx 0, partial_sum[0] 14710790.000000
threadIdx 1, partial_sum[1] 18872320.000000
threadIdx 0, partial_sum[0] 33583112.000000
threadIdx 1, partial_sum[1] 33059334.000000
expected 33583110, result 33583112

前两行部分和相加14710790+18872320=33583110是正确值，但最终求和结果错误。

问题根源 #

问题出在warp级无同步的归约步骤中：

• 当执行partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 1]时，thread 0和thread 1会同时读写共享内存，导致数据竞争。
• thread 0读取partial_sum[1]的同时，thread 1也在更新自己的partial_sum[1]（thread 1会执行partial_sum[1] += partial_sum[2]），导致thread 0拿到的是thread 1更新后的值，而非初始的18872320。

修复方案 #

在warp级归约的最后一步，只有thread 0需要执行跨thread的加法，其他thread不需要参与最终的全局归约。修改warp内的归约逻辑，确保只有thread 0完成最终的+1步：

if (threadIdx.x < 32) {
    partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 32];
    partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 16];
    partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 8];
    partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 4];
    partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 2];
    if (threadIdx.x < 2 && blockIdx.x == 0 && print) {
        printf("threadIdx %d, partial_sum[%d] %f\n", threadIdx.x, threadIdx.x, partial_sum[threadIdx.x]);
    }

    // 只有thread 0执行最终的+1操作，避免数据竞争
    if (threadIdx.x == 0) {
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 1];
    }

    if (threadIdx.x < 2 && blockIdx.x == 0 && print) {
        printf("threadIdx %d, partial_sum[%d] %f\n", threadIdx.x, threadIdx.x, partial_sum[threadIdx.x]);
    }
}

另外补充两个优化点：

1. 共享内存分配冗余：当前分配了THREADS_PER_DIM*4的共享内存，但实际只用到THREADS_PER_DIM，可改为volatile __shared__ float partial_sum[THREADS_PER_DIM];避免浪费。
2. num_blocks计算不通用：手动指定num_blocks=9不具备通用性，改用公式动态计算：

   auto num_blocks = (N + 2*THREADS_PER_DIM - 1) / (2*THREADS_PER_DIM);
   

完整修复后的代码 #

#include <iostream>
#include <numeric>
#include <vector>

constexpr unsigned int THREADS_PER_DIM = 512;

__global__ void reduce(float *v, float *v_r, std::size_t N, bool print) {
    // 优化共享内存大小
    volatile __shared__ float partial_sum[THREADS_PER_DIM];

    unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;

    partial_sum[threadIdx.x] = 0;
    partial_sum[threadIdx.x] = (i < N ? v[i] : 0) + ((i + blockDim.x) < N ? v[i + blockDim.x] : 0);
    __syncthreads();

    for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {
        if (threadIdx.x < s) {
            partial_sum[threadIdx.x] += partial_sum[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x < 32) {
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 32];
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 16];
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 8];
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 4];
        partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 2];
        if (threadIdx.x < 2 && blockIdx.x == 0 && print) {
            printf("threadIdx %d, partial_sum[%d] %f\n", threadIdx.x, threadIdx.x, partial_sum[threadIdx.x]);
        }

        // 仅thread 0执行最终的+1操作
        if (threadIdx.x == 0) {
            partial_sum[threadIdx.x] += partial_sum[threadIdx.x + 1];
        }

        if (threadIdx.x < 2 && blockIdx.x == 0 && print) {
            printf("threadIdx %d, partial_sum[%d] %f\n", threadIdx.x, threadIdx.x, partial_sum[threadIdx.x]);
        }
    }

    if (threadIdx.x == 0) {
        v_r[blockIdx.x] = partial_sum[0];
    }
}

constexpr std::size_t N = (1 << 13) + 4;
int main() {
    std::vector<float> a(N);
    for (std::size_t i = 0; i < N; ++i) {
        a[i] = static_cast<float>(i);
    }

    std::size_t bytes = N * sizeof(float);
    float *d_A;
    float *d_A_reduction;
    cudaMalloc((void **)&d_A, bytes);
    cudaMalloc((void **)&d_A_reduction, bytes);
    cudaMemset(d_A_reduction, 0, bytes);
    cudaMemcpy(d_A, a.data(), bytes, cudaMemcpyHostToDevice);

    // 动态计算block数量
    auto num_blocks = (N + 2*THREADS_PER_DIM - 1) / (2*THREADS_PER_DIM);

    reduce<<<num_blocks, THREADS_PER_DIM>>>(d_A, d_A_reduction, N, false);
    cudaMemcpy(d_A, d_A_reduction, bytes, cudaMemcpyDeviceToDevice);
    reduce<<<1, THREADS_PER_DIM>>>(d_A, d_A_reduction, num_blocks, true);

    float result = 0;
    cudaMemcpy(&result, d_A_reduction, sizeof(float), cudaMemcpyHostToDevice);
    cudaFree(d_A);
    cudaFree(d_A_reduction);

    auto correct_result = std::accumulate(a.begin(), a.end(), 0.0);
    std::cout << "expected " << static_cast<int>(correct_result) << ", result " << static_cast<int>(result)
              << std::endl;
}

编译运行nvcc test.cu -o test && ./test，即可得到正确结果。

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