如何使用CUDA Graphs，如何更新Graphs中kernel函数参数

使用CUDA Graphs及更新Kernel参数示例

CUDA Graphs提供了一种高效的方式来捕获和执行一系列CUDA操作，特别适合重复执行相同操作序列的场景。下面我将展示如何使用CUDA Graphs以及如何更新Graph中的kernel参数。

基本CUDA Graphs使用示例

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

#define N 1024

__global__ void addKernel(int *a, int *b, int *c, int scale) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        c[i] = a[i] + b[i] * scale;
    }
}

int main() {
    // 初始化主机数据
    int ha[N], hb[N], hc[N];
    for (int i = 0; i < N; i++) {
        ha[i] = i;
        hb[i] = i * 2;
    }
    
    // 分配设备内存
    int *da, *db, *dc;
    cudaMalloc(&da, N * sizeof(int));
    cudaMalloc(&db, N * sizeof(int));
    cudaMalloc(&dc, N * sizeof(int));
    
    // 拷贝数据到设备
    cudaMemcpy(da, ha, N * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(db, hb, N * sizeof(int), cudaMemcpyHostToDevice);
    
    // 创建CUDA Graph
    cudaGraph_t graph;
    cudaGraphExec_t graphExec;
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    
    // 开始捕获Graph
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
    
    // 在流中执行操作(这些操作将被捕获到Graph中)
    dim3 block(256);
    dim3 grid((N + block.x - 1) / block.x);
    int scale = 1; // 初始参数值
    addKernel<<<grid, block, 0, stream>>>(da, db, dc, scale);
    
    // 结束捕获并实例化Graph
    cudaStreamEndCapture(stream, &graph);
    cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0);
    
    // 执行Graph
    cudaGraphLaunch(graphExec, stream);
    cudaStreamSynchronize(stream);
    
    // 检查结果
    cudaMemcpy(hc, dc, N * sizeof(int), cudaMemcpyDeviceToHost);
    printf("Result[0] = %d (should be 0)\n", hc[0]);
    printf("Result[1] = %d (should be 3)\n", hc[1]);
    
    // 更新Graph中的kernel参数并重新执行
    // 方法1: 重新捕获Graph(适用于参数或结构变化较大的情况)
    cudaGraphDestroy(graph);
    cudaGraphExecDestroy(graphExec);
    
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
    scale = 2; // 更新参数值
    addKernel<<<grid, block, 0, stream>>>(da, db, dc, scale);
    cudaStreamEndCapture(stream, &graph);
    cudaGraphInstantiate(&graphExec, graph, NULL, NULL, 0);
    
    cudaGraphLaunch(graphExec, stream);
    cudaStreamSynchronize(stream);
    
    cudaMemcpy(hc, dc, N * sizeof(int), cudaMemcpyDeviceToHost);
    printf("After update:\n");
    printf("Result[0] = %d (should be 0)\n", hc[0]);
    printf("Result[1] = %d (should be 5)\n", hc[1]);
    
    // 方法2: 使用cudaGraphExecKernelNodeSetParams更新单个kernel参数(更高效)
    // 首先找到kernel节点
    cudaGraphNode_t* nodes = NULL;
    size_t numNodes = 0;
    cudaGraphGetNodes(graph, nodes, &numNodes); // 第一次获取节点数量
    nodes = (cudaGraphNode_t*)malloc(numNodes * sizeof(cudaGraphNode_t));
    cudaGraphGetNodes(graph, nodes, numNodes); // 第二次获取节点
    
    // 遍历节点找到kernel节点
    cudaGraphNode_t kernelNode = NULL;
    for (size_t i = 0; i < numNodes; i++) {
        cudaGraphNodeType type;
        cudaGraphNodeGetType(nodes[i], &type);
        if (type == cudaGraphNodeTypeKernel) {
            kernelNode = nodes[i];
            break;
        }
    }
    
    if (kernelNode) {
        // 设置新的kernel参数
        scale = 3;
        cudaKernelNodeParams params;
        memset(&params, 0, sizeof(params));
        params.func = (void*)addKernel;
        params.gridDim = grid;
        params.blockDim = block;
        params.sharedMemBytes = 0;
        params.kernelParams = (void**)&da, &db, &dc, &scale;
        params.extra = NULL;
        
        // 更新参数
        cudaGraphExecKernelNodeSetParams(graphExec, kernelNode, &params);
        
        // 再次执行
        cudaGraphLaunch(graphExec, stream);
        cudaStreamSynchronize(stream);
        
        cudaMemcpy(hc, dc, N * sizeof(int), cudaMemcpyDeviceToHost);
        printf("After cudaGraphExecKernelNodeSetParams:\n");
        printf("Result[0] = %d (should be 0)\n", hc[0]);
        printf("Result[1] = %d (should be 7)\n", hc[1]);
    }
    
    // 清理资源
    free(nodes);
    cudaGraphExecDestroy(graphExec);
    cudaGraphDestroy(graph);
    cudaStreamDestroy(stream);
    cudaFree(da);
    cudaFree(db);
    cudaFree(dc);
    
    return 0;
}

关键点说明

1. Graph捕获:

   • 使用cudaStreamBeginCapture和cudaStreamEndCapture捕获操作序列
   • 捕获期间在流上执行的所有操作都会被记录到Graph中

2. Graph执行:

   • 使用cudaGraphInstantiate创建可执行Graph实例
   • 使用cudaGraphLaunch执行Graph

3. 参数更新方法:

   • 方法1: 重新捕获整个Graph - 简单但效率较低
   • 方法2: 使用cudaGraphExecKernelNodeSetParams直接更新kernel参数 - 更高效但代码更复杂

4. 性能考虑:

   • CUDA Graphs最适合操作序列固定但参数可能变化的场景
   • 对于频繁变化的参数，使用方法2更新更高效
   • 如果操作结构变化很大，可能需要重新捕获Graph

编译运行

使用nvcc编译:

nvcc cuda_graph_example.cu -o cuda_graph_example
./cuda_graph_example

这个示例展示了CUDA Graphs的基本用法和参数更新技术，你可以根据实际需求调整kernel函数和参数更新逻辑。