CUDA:使用线程束洗牌函数进行规约计算

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

typedef float real; // 你可以根据需要修改数据类型

const int N = 1024; // 输入数组大小
const int M = sizeof(real) * N;
const int BLOCK_SIZE = 128;
const unsigned FULL_MASK = 0xffffffff;
// CUDA核函数
__global__ void reduce_shfl(const real* d_x, real* d_y, const int N)
{
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int n = bid * blockDim.x + tid;
    extern __shared__ real s_y[];
    s_y[tid] = (n < N) ? d_x[n] : 0.0;
    __syncthreads();

    for (int offset = blockDim.x >> 1; offset >= 32; offset >>= 1)
    {
        if (tid < offset)
        {
            s_y[tid] += s_y[tid + offset];
        }
        __syncthreads();
    }

    real y = s_y[tid];

    for (int offset = 16; offset > 0; offset >>= 1)
    {
        y += __shfl_down_sync(FULL_MASK, y, offset);
    }

    if (tid == 0)
    {
        atomicAdd(d_y, y);
    }
}

int main()
{
    real* d_x, * d_y;
    real h_x[N]; // 输入数组
    real h_y = 0.0; // 输出结果

    // 初始化输入数组
    for (int i = 0; i < N; i++)
    {
        h_x[i] = 2; // 你可以根据需要修改初始化值
    }

    // 在GPU上分配内存
    cudaMalloc((void**)&d_x, N * sizeof(real));
    cudaMalloc((void**)&d_y, sizeof(real));

    // 将输入数组复制到GPU
    cudaMemcpy(d_x, h_x, N * sizeof(real), cudaMemcpyHostToDevice);
    cudaMemcpy(d_y, &h_y, sizeof(real), cudaMemcpyHostToDevice);

    // 设置线程块大小和网格大小
    int blockSize = 256;
    int gridSize = (N + blockSize - 1) / blockSize;

    // 调用核函数
    reduce_shfl << <gridSize, blockSize, blockSize * sizeof(real) >> > (d_x, d_y, N);

    // 将结果从GPU复制回主机
    cudaMemcpy(&h_y, d_y, sizeof(real), cudaMemcpyDeviceToHost);

    // 打印结果
    std::cout << "Result: " << h_y << std::endl;

    // 释放GPU内存
    cudaFree(d_x);
    cudaFree(d_y);

    return 0;
}