CUDA Graph Demo

背景

假设我们有一个深度学习的训练过程，其中包含了多个计算操作，比如：

从主机内存（CPU）传输数据到设备内存（GPU）。
执行 GPU 上的计算内核（例如，前向传播、反向传播）。
从设备内存传输计算结果回主机内存。

如果每次都单独执行这些操作，CUDA 就需要多次调度这些操作并进行同步，这会增加 CPU 和 GPU 之间的通信开销。

传统方式（没有 CUDA Graph）

在传统的 CUDA 编程中，你可能会像这样手动管理每个步骤：

#include <cuda_runtime.h>

__global__ void kernel_forward(float* data) {
    // 假设这是一个前向传播的计算内核
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    data[idx] = data[idx] * 2.0f;  // 简单的计算
}

__global__ void kernel_backward(float* data) {
    // 假设这是一个反向传播的计算内核
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    data[idx] = data[idx] / 2.0f;  // 简单的计算
}

int main() {
    float *d_data;
    float *h_data = (float*)malloc(sizeof(float) * 1024);

    // Step 1: 数据从主机传输到设备
    cudaMalloc(&d_data, sizeof(float) * 1024);
    cudaMemcpy(d_data, h_data, sizeof(float) * 1024, cudaMemcpyHostToDevice);

    // Step 2: 执行前向传播内核
    kernel_forward<<<1, 1024>>>(d_data);
    cudaDeviceSynchronize();

    // Step 3: 执行反向传播内核
    kernel_backward<<<1, 1024>>>(d_data);
    cudaDeviceSynchronize();

    // Step 4: 数据从设备传输回主机
    cudaMemcpy(h_data, d_data, sizeof(float) * 1024, cudaMemcpyDeviceToHost);

    cudaFree(d_data);
    free(h_data);

    return 0;
}

在这个传统的例子中，CPU 会逐步执行数据传输和内核调用。在每个内核执行后，cudaDeviceSynchronize() 需要确保 GPU 完成当前操作后才能继续执行下一步。这会造成不必要的延迟和调度开销。

使用 CUDA Graph 的优化

通过使用 CUDA Graph，我们可以捕获和优化这些操作的执行顺序，减少每次执行时的调度开销。下面是一个简单的改进版，使用 CUDA Graph 来捕获操作并执行它们。

#include <cuda_runtime.h>
#include <cuda_runtime_api.h>

__global__ void kernel_forward(float* data) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    data[idx] = data[idx] * 2.0f;  // 简单的计算
}

__global__ void kernel_backward(float* data) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    data[idx] = data[idx] / 2.0f;  // 简单的计算
}

int main() {
    float *d_data;
    float *h_data = (float*)malloc(sizeof(float) * 1024);

    // Step 1: 数据从主机传输到设备
    cudaMalloc(&d_data, sizeof(float) * 1024);
    cudaMemcpy(d_data, h_data, sizeof(float) * 1024, cudaMemcpyHostToDevice);

    // Step 2: 创建 CUDA Graph
    cudaGraph_t graph;
    cudaGraphCreate(&graph, 0);

    // Step 3: 捕获内核节点
    cudaKernelNodeParams forward_params = {};
    forward_params.func = (void*)kernel_forward;
    forward_params.gridDim = dim3(1);
    forward_params.blockDim = dim3(1024);
    forward_params.kernelParams = (void**)&d_data;

    cudaKernelNodeParams backward_params = {};
    backward_params.func = (void*)kernel_backward;
    backward_params.gridDim = dim3(1);
    backward_params.blockDim = dim3(1024);
    backward_params.kernelParams = (void**)&d_data;

    // 创建 CUDA 图形节点
    cudaGraphNode_t forward_node, backward_node;
    cudaGraphAddKernelNode(&forward_node, graph, nullptr, 0, &forward_params);
    cudaGraphAddKernelNode(&backward_node, graph, &forward_node, 1, &backward_params);

    // Step 4: 提交图形执行
    cudaGraphLaunch(graph, 0);
    cudaDeviceSynchronize();  // 等待执行完成

    // Step 5: 数据从设备传输回主机
    cudaMemcpy(h_data, d_data, sizeof(float) * 1024, cudaMemcpyDeviceToHost);

    cudaFree(d_data);
    free(h_data);
    cudaGraphDestroy(graph);

    return 0;
}

解释

捕获操作：
- 我们使用 cudaGraphAddKernelNode 捕获了两个内核（前向传播和反向传播）以及它们的依赖关系。在这个图形中，前向传播内核必须在反向传播内核之前执行，反向传播内核是依赖于前向传播的结果的。
优化执行：
- 捕获完图形后，cudaGraphLaunch 会提交图形中的所有操作，而不需要每次都单独调用每个内核，并且 CUDA 会根据依赖关系进行优化。在这种情况下，图形中的操作已经明确了执行顺序，GPU 可以按图形的顺序高效地执行任务。
减少调度和通信开销：
- 在图形执行时，GPU 不需要每次都等待 CPU 调度。所有的操作和依赖关系已经在图形中捕获，GPU 只需按照预定的顺序执行，无需每次都进行调度。
- 同时，cudaMemcpy 操作也可以作为图形的一部分被捕获，减少了 CPU 和 GPU 之间的频繁交互。

总结

在这个示例中，CUDA Graph 通过将多个操作（如数据传输和计算内核）捕获为一个图形，减少了每次执行时的调度开销和 CPU 与 GPU 之间的通信开销。通过预先定义好操作的依赖关系，GPU 可以更高效地执行这些操作，避免了重复的任务调度和同步操作，从而显著提高了性能。