Softmax CUDA实现

// Softmax: P(x_i) = exp(x_i) / Σexp(x_j)
//
// 数值稳定版本:
// 1. 找最大值 m = max(x_i)
// 2. 计算 exp(x_i - m)  (减去最大值防止溢出)
// 3. 求和 s = Σexp(x_i - m)
// 4. 归一化 P(x_i) = exp(x_i - m) / s

/**
 * Softmax CUDA实现 - 面试手撕版
 * 
 * 特点:
 * - 数值稳定 (减去最大值)
 * - 支持大数组 (多Block协作)
 * - 使用Warp Shuffle高效归约
 * 
 * 编译: nvcc -o softmax softmax.cu
 */

#include <stdio.h>
#include <math.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归约 ====================
__device__ float warp_reduce_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;
}

__device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xFFFFFFFF, val, offset);
    }
    return val;
}

// ==================== Softmax Kernel ====================
__global__ void softmax_kernel(const float* input, float* output, int n) {
    extern __shared__ float smem[];
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    // Step 1: 加载数据
    float val = (i < n) ? input[i] : -INFINITY;
    
    // Step 2: Warp内找最大值
    float max_val = warp_reduce_max(val);
    
    // 每个Warp的最大值写入Shared Memory
    if (lane_id == 0) {
        smem[warp_id] = max_val;
    }
    __syncthreads();
    
    // Step 3: 第一个Warp归约所有Warp的最大值
    if (warp_id == 0) {
        max_val = (tid < blockDim.x / 32) ? smem[tid] : -INFINITY;
        max_val = warp_reduce_max(max_val);
        if (tid == 0) smem[0] = max_val;  // 全局最大值
    }
    __syncthreads();
    float global_max = smem[0];
    
    // Step 4: 计算exp(x - max)和求和
    val = (i < n) ? expf(input[i] - global_max) : 0.0f;
    
    // Warp内求和
    float sum_val = warp_reduce_sum(val);
    
    if (lane_id == 0) {
        smem[warp_id] = sum_val;
    }
    __syncthreads();
    
    // Step 5: 第一个Warp归约所有Warp的和
    if (warp_id == 0) {
        sum_val = (tid < blockDim.x / 32) ? smem[tid] : 0.0f;
        sum_val = warp_reduce_sum(sum_val);
        if (tid == 0) smem[0] = sum_val;  // 全局和
    }
    __syncthreads();
    float global_sum = smem[0];
    
    // Step 6: 归一化
    if (i < n) {
        output[i] = val / global_sum;
    }
}

// ==================== CPU参考实现 ====================
void softmax_cpu(const float* input, float* output, int n) {
    float max_val = *std::max_element(input, input + n);
    float sum = 0;
    for (int i = 0; i < n; i++) {
        output[i] = expf(input[i] - max_val);
        sum += output[i];
    }
    for (int i = 0; i < n; i++) {
        output[i] /= sum;
    }
}

// ==================== 主函数 ====================
int main() {
    const int N = 1024;
    const int block_size = 256;
    
    float* h_input = (float*)malloc(N * sizeof(float));
    float* h_output_gpu = (float*)malloc(N * sizeof(float));
    float* h_output_cpu = (float*)malloc(N * sizeof(float));
    
    // 初始化 (随机值)
    for (int i = 0; i < N; i++) {
        h_input[i] = (float)(rand() % 100) / 100.0f;
    }
    
    // GPU计算
    float *d_input, *d_output;
    CUDA_CHECK(cudaMalloc(&d_input, N * sizeof(float)));
    CUDA_CHECK(cudaMalloc(&d_output, N * sizeof(float)));
    CUDA_CHECK(cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice));
    
    softmax_kernel<<<>N/block_size, block_size, block_size * sizeof(float)>>>(
        d_input, d_output, N);
    CUDA_CHECK(cudaDeviceSynchronize());
    
    CUDA_CHECK(cudaMemcpy(h_output_gpu, d_output, N * sizeof(float), cudaMemcpyDeviceToHost));
    
    // CPU验证
    softmax_cpu(h_input, h_output_cpu, N);
    
    // 对比
    float max_diff = 0;
    for (int i = 0; i < N; i++) {
        max_diff = fmaxf(max_diff, fabs(h_output_gpu[i] - h_output_cpu[i]));
    }
    printf("Max difference: %e\n", max_diff);
    printf(max_diff < 1e-5 ? "✅ PASS\n" : "❌ FAIL\n");
    
    // 验证概率和为1
    float sum = 0;
    for (int i = 0; i < N; i++) sum += h_output_gpu[i];
    printf("Sum of probabilities: %.6f\n", sum);
    
    // 清理
    cudaFree(d_input); cudaFree(d_output);
    free(h_input); free(h_output_gpu); free(h_output_cpu);
    
    return 0;
}