Stream并发Pipeline

/**
 * Stream并发Pipeline实现
 * 
 * 对比:
 * 1. 串行执行（默认Stream）
 * 2. 多Stream并发（计算传输重叠）
 * 
 * 编译: nvcc -o pipeline pipeline.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)

// 简单的计算Kernel
__global__ void compute_kernel(float* data, int n, float factor) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < n) {
        // 模拟计算密集型操作
        float val = data[idx];
        for (int i = 0; i < 100; i++) {
            val = val * factor + 1.0f;
        }
        data[idx] = val;
    }
}

// ==================== 串行执行 ====================
void serial_execution(float* h_data, float* d_data, int n, int num_chunks) {
    int chunk_size = n / num_chunks;
    int chunk_bytes = chunk_size * sizeof(float);
    
    for (int i = 0; i < num_chunks; i++) {
        // Host → Device
        CUDA_CHECK(cudaMemcpy(d_data + i * chunk_size, 
                              h_data + i * chunk_size, 
                              chunk_bytes, cudaMemcpyHostToDevice));
        
        // 计算
        compute_kernel<<<chunk_size/256, 256>>>(d_data + i * chunk_size, chunk_size, 0.99f);
        
        // Device → Host
        CUDA_CHECK(cudaMemcpy(h_data + i * chunk_size,
                              d_data + i * chunk_size,
                              chunk_bytes, cudaMemcpyDeviceToHost));
    }
}

// ==================== Stream并发执行 ====================
void stream_execution(float* h_data, float* d_data, int n, int num_streams) {
    int stream_size = n / num_streams;
    int stream_bytes = stream_size * sizeof(float);
    
    // 创建Streams
    cudaStream_t* streams = new cudaStream_t[num_streams];
    for (int i = 0; i < num_streams; i++) {
        cudaStreamCreate(&streams[i]);
    }
    
    // 为每个Stream分配Host pinned memory（可选，提高传输速度）
    float** h_chunk = new float*[num_streams];
    for (int i = 0; i < num_streams; i++) {
        cudaMallocHost(&h_chunk[i], stream_bytes);  // Pinned memory
        // 复制数据
        for (int j = 0; j < stream_size; j++) {
            h_chunk[i][j] = h_data[i * stream_size + j];
        }
    }
    
    // 并发执行
    for (int i = 0; i < num_streams; i++) {
        // Host → Device (异步)
        CUDA_CHECK(cudaMemcpyAsync(d_data + i * stream_size,
                                    h_chunk[i],
                                    stream_bytes,
                                    cudaMemcpyHostToDevice,
                                    streams[i]));
        
        // 计算（在指定Stream中执行）
        compute_kernel<<<>stream_size/256, 256, 0, streams[i]>>>(
            d_data + i * stream_size, stream_size, 0.99f);
        
        // Device → Host (异步)
        CUDA_CHECK(cudaMemcpyAsync(h_chunk[i],
                                    d_data + i * stream_size,
                                    stream_bytes,
                                    cudaMemcpyDeviceToHost,
                                    streams[i]));
    }
    
    // 等待所有Stream完成
    for (int i = 0; i < num_streams; i++) {
        CUDA_CHECK(cudaStreamSynchronize(streams[i]));
    }
    
    // 清理
    for (int i = 0; i < num_streams; i++) {
        cudaStreamDestroy(streams[i]);
        cudaFreeHost(h_chunk[i]);
    }
    delete[] streams;
    delete[] h_chunk;
}

// ==================== 性能测试 ====================
int main() {
    const int N = 1 << 22;  // 4M elements
    const int num_chunks = 4;
    
    printf("Stream并发Pipeline演示\n");
    printf("数据大小: %.2f MB\n", N * sizeof(float) / 1024.0 / 1024.0);
    
    // 分配内存
    float *h_data = (float*)malloc(N * sizeof(float));
    float *d_data;
    CUDA_CHECK(cudaMalloc(&d_data, N * sizeof(float)));
    
    // 初始化
    for (int i = 0; i < N; i++) h_data[i] = float(i);
    
    // 计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start); cudaEventCreate(&stop);
    float ms;
    
    // 串行执行
    cudaEventRecord(start);
    serial_execution(h_data, d_data, N, num_chunks);
    cudaEventRecord(stop); cudaEventSynchronize(stop);
    cudaEventElapsedTime(&ms, start, stop);
    printf("串行执行: %.2f ms\n", ms);
    
    // Stream并发
    cudaEventRecord(start);
    stream_execution(h_data, d_data, N, num_chunks);
    cudaEventRecord(stop); cudaEventSynchronize(stop);
    cudaEventElapsedTime(&ms, start, stop);
    printf("Stream并发: %.2f ms\n", ms);
    
    // 清理
    cudaFree(d_data); free(h_data);
    
    return 0;
}