Warp Shuffle归约

/**
 * Warp Shuffle归约详解
 * 
 * 演示:
 * 1. 基本Warp Shuffle操作
 * 2. Warp内求和归约
 * 3. Warp内求最大值
 * 4. 多Block协作归约
 * 
 * 编译: nvcc -o warp_shuffle warp_shuffle.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)

// ==================== 基本Warp Shuffle操作 ====================
// __shfl_sync: 从指定lane获取数据
// __shfl_up_sync: 从上方lane获取数据
// __shfl_down_sync: 从下方lane获取数据
// __shfl_xor_sync: 从异或位置获取数据

// ==================== Warp求和归约 ====================
__device__ float warp_sum(float val) {
    // 方法1: 使用__shfl_down_sync（最常用）
    // 每次迭代，将下方offset个线程的值加到当前线程
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xFFFFFFFF, val, offset);
    }
    return val;
}

// 方法2: 使用__shfl_xor_sync（蝶形归约）
__device__ float warp_sum_xor(float val) {
    // XOR掩码: 1, 2, 4, 8, 16
    for (int offset = 1; offset < 32; offset <<= 1) {
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    }
    return val;
}

// ==================== Warp最大值归约 ====================
__device__ float warp_max(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xFFFFFFFF, val, offset);
        val = fmaxf(val, other);
    }
    return val;
}

// ==================== 多Block协作归约 ====================
__global__ void reduce_sum(const float* input, float* output, int n) {
    extern __shared__ float smem[];
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + tid;
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    // 加载数据
    float val = (i < n) ? input[i] : 0.0f;
    
    // Step 1: Warp内归约（32个线程用5步）
    val = warp_sum(val);
    
    // Step 2: 每个Warp的线程0写入Shared Memory
    if (lane_id == 0) {
        smem[warp_id] = val;
    }
    __syncthreads();
    
    // Step 3: 第一个Warp归约所有Warp的结果
    if (warp_id == 0) {
        // 只有前num_warps个线程参与
        int num_warps = blockDim.x / 32;
        val = (tid < num_warps) ? smem[tid] : 0.0f;
        val = warp_sum(val);
    }
    
    // Step 4: 线程0写入最终结果
    if (tid == 0) {
        output[blockIdx.x] = val;
    }
}

// ==================== 测试Warp Shuffle基本操作 ====================
__global__ void test_shuffle(float* output) {
    int tid = threadIdx.x;
    float val = (float)tid;
    
    // 测试__shfl_down_sync
    float down_val = __shfl_down_sync(0xFFFFFFFF, val, 1);
    // 线程i的down_val = 线程i+1的val
    
    // 测试__shfl_xor_sync
    float xor_val = __shfl_xor_sync(0xFFFFFFFF, val, 1);
    // 线程i的xor_val = 线程i^1的val
    
    output[tid] = down_val + xor_val;
}

// ==================== 主函数 ====================
int main() {
    const int N = 1024;
    const int block_size = 256;
    const int grid_size = (N + block_size - 1) / block_size;
    
    printf("Warp Shuffle归约演示\n");
    printf("Warp大小: 32\n");
    printf("Block大小: %d (%d warps)\n", block_size, block_size / 32);
    
    // 测试1: 基本Shuffle操作
    printf("\n=== 测试1: 基本Shuffle操作 ===\n");
    float *d_out;
    CUDA_CHECK(cudaMalloc(&d_out, 32 * sizeof(float)));
    test_shuffle<<<1, 32>>>(d_out);
    
    float h_out[32];
    CUDA_CHECK(cudaMemcpy(h_out, d_out, 32 * sizeof(float), cudaMemcpyDeviceToHost));
    printf("Thread 0 result: %.1f (expected: 0+1 + 0^1=1 = 1.0)\n", h_out[0]);
    printf("Thread 5 result: %.1f (expected: 6 + 5^1=4 = 10.0)\n", h_out[5]);
    
    // 测试2: 向量归约
    printf("\n=== 测试2: 向量归约求和 ===\n");
    
    float* h_input = (float*)malloc(N * sizeof(float));
    for (int i = 0; i < N; i++) h_input[i] = 1.0f;
    
    float *d_input;
    CUDA_CHECK(cudaMalloc(&d_input, N * sizeof(float)));
    CUDA_CHECK(cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));
    
    reduce_sum<<<grid_size, block_size, block_size * sizeof(float)>>>(
        d_input, d_out, N);
    
    CUDA_CHECK(cudaMemcpy(h_out, d_out, grid_size * sizeof(float), cudaMemcpyDeviceToHost));
    
    float total = 0;
    for (int i = 0; i < grid_size; i++) total += h_out[i];
    printf("Sum: %.0f (expected: %d)\n", total, N);
    printf(total == N ? "✅ PASS\n" : "❌ FAIL\n");
    
    // 清理
    cudaFree(d_input); cudaFree(d_out);
    free(h_input);
    
    return 0;
}