向量归约求和

/**
 * 向量归约求和 - 三种方法对比
 * 
 * 方法1: 朴素归约 (O(n) 串行)
 * 方法2: Shared Memory归约 (O(log n) 并行)
 * 方法3: Warp Shuffle归约 (最快，无共享内存开销)
 * 
 * 编译: nvcc -o reduction reduction.cu
 */

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

#define CUDA_CHECK(call) do { cudaError_t e = call; if(e != cudaSuccess) { \
    fprintf(stderr, "CUDA Error: %s\n", cudaGetErrorString(e)); exit(1); } } while(0)

// ==================== 方法1: 朴素归约 ====================
// 每个Block只用一个线程串行求和
__global__ void reduce_naive(const float* input, float* output, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    
    // 只有每个Block的第一个线程执行
    if (threadIdx.x == 0) {
        float sum = 0.0f;
        for (int i = blockIdx.x * blockDim.x; 
             i < min((int)(blockIdx.x * blockDim.x + blockDim.x), n); 
             i++) {
            sum += input[i];
        }
        output[blockIdx.x] = sum;
    }
}

// ==================== 方法2: Shared Memory归约 ====================
// 树形归约，O(log n)复杂度
__global__ void reduce_shared(const float* input, float* output, int n) {
    extern __shared__ float sdata[];
    
    // 每个线程加载一个元素到Shared Memory
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    sdata[tid] = (i < n) ? input[i] : 0.0f;
    __syncthreads();
    
    // 树形归约
    // 每次迭代，活跃线程数减半
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // 只有线程0写入结果
    if (tid == 0) {
        output[blockIdx.x] = sdata[0];
    }
}

// ==================== 方法3: Warp Shuffle归约 ====================
// 利用Warp内线程直接交换数据，最快
__device__ float warp_reduce_sum(float val) {
    // __shfl_down_sync: 将当前线程的值发送给下方offset个线程
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xFFFFFFFF, val, offset);
    }
    return val;
}

__global__ void reduce_warp_shuffle(const float* input, float* output, int n) {
    extern __shared__ float sdata[];
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 加载数据
    float val = (i < n) ? input[i] : 0.0f;
    
    // Warp内归约（32个线程用5步）
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    // 先做Warp内归约
    val = warp_reduce_sum(val);
    
    // 每个Warp的结果写入Shared Memory
    if (lane_id == 0) {
        sdata[warp_id] = val;
    }
    __syncthreads();
    
    // 第一个Warp归约所有Warp的结果
    if (warp_id == 0) {
        val = (tid < blockDim.x / 32) ? sdata[tid] : 0.0f;
        val = warp_reduce_sum(val);
    }
    
    if (tid == 0) {
        output[blockIdx.x] = val;
    }
}

// ==================== 主函数 ====================
int main() {
    const int N = 1 << 20;  // 1M elements
    const int block_size = 256;
    const int grid_size = (N + block_size - 1) / block_size;
    
    float *h_input = (float*)malloc(N * sizeof(float));
    float *h_output = (float*)malloc(grid_size * sizeof(float));
    
    // 初始化
    for (int i = 0; i < N; i++) h_input[i] = 1.0f;
    
    float *d_input, *d_output;
    CUDA_CHECK(cudaMalloc(&d_input, N * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_output, grid_size * sizeof(float)));
    CUDA_CHECK(cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));
    
    // 测试三种方法
    printf("向量归约求和: N = %d\n", N);
    
    // 方法2: Shared Memory
    reduce_shared<<<grid_size, block_size, block_size * sizeof(float)>>>(
        d_input, d_output, N);
    CUDA_CHECK(cudaDeviceSynchronize());
    
    // 结果验证
    CUDA_CHECK(cudaMemcpy(h_output, d_output, grid_size * sizeof(float), cudaMemcpyDeviceToHost));
    float sum = 0;
    for (int i = 0; i < grid_size; i++) sum += h_output[i];
    printf("Shared Memory Result: %.0f (expected %d)\n", sum, N);
    
    // 清理
    cudaFree(d_input); cudaFree(d_output);
    free(h_input); free(h_output);
    
    return 0;
}