CUDA Graph运行结果不符合预期，求技术解答

问题描述 #

我正在用以下代码学习CUDA Graph，设置NSTEP=1000、NKERNEL=20。核函数shortKernel执行简单计算，但实际运行结果完全不符合预期。

原代码 #

#include <cuda_runtime.h>
#include <iostream>

#define N 131072 // tuned such that kernel takes a few microseconds
#define NSTEP 1000
#define NKERNEL 20
#define BLOCKS 256
#define THREADS 512

#define CHECK(call)                                                         \
    do {                                                                    \
        const cudaError_t error_code = call;                                \
        if (error_code != cudaSuccess) {                                    \
            printf("CUDA Error\n");                                         \
            printf("    File:   %s\n", __FILE__);                           \
            printf("    Line:   %d\n", __LINE__);                           \
            printf("    Error code: %d\n", error_code);                     \
            printf("    Error text: %s\n", cudaGetErrorString(error_code)); \
            exit(1);                                                        \
        }                                                                   \
    } while (0)

__global__ void shortKernel(float * out_d, float * in_d, int i){
      int idx=blockIdx.x*blockDim.x+threadIdx.x;
        if(idx<N) out_d[idx]=1.23*in_d[idx] + i;
        
}

void test2() {
  cudaStream_t stream;
  cudaStreamCreate(&stream);
  cudaSetDevice(0);

  float x_host[N], y_host[N];
  // initialize x and y arrays on the host
  for (int i = 0; i < N; i++) {
    x_host[i] = 2.0f;
    y_host[i] = 2.0f;
  }
  float *x, *y, *z;
  CHECK(cudaMalloc((void**)&x, N*sizeof(float)));
  CHECK(cudaMalloc((void**)&y, N*sizeof(float)));
  CHECK(cudaMalloc((void**)&z, N*sizeof(float)));
  cudaMemcpy(x, x_host, sizeof(float) * N, cudaMemcpyHostToDevice);

  cudaEvent_t begin, end;
  CHECK(cudaEventCreate(&begin));
  CHECK(cudaEventCreate(&end));
  // start recording
  cudaEventRecord(begin, stream);
  bool graphCreated=false;
  cudaGraph_t graph;
  cudaGraphExec_t instance;
  // Run graphs
  for(int istep=0; istep<NSTEP; istep++){
    if(!graphCreated){
      cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
      for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
        shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, ikrnl);
      }
      cudaStreamEndCapture(stream, &graph);
      cudaGraphNode_t* nodes = NULL;
      size_t num_nodes = 0;
      CHECK(cudaGraphGetNodes(graph, nodes, &num_nodes));
      std::cout << "Num of nodes in the graph: " << num_nodes
                << std::endl;
      CHECK(cudaGraphInstantiate(&instance, graph, NULL, NULL, 0));
      graphCreated=true;
    }
    CHECK(cudaGraphLaunch(instance, stream));
    cudaStreamSynchronize(stream);
  }  // End run graphs
  cudaEventRecord(end, stream);
  cudaEventSynchronize(end);
  float time_ms = 0;
  cudaEventElapsedTime(&time_ms, begin, end);
  std::cout << "CUDA Graph - CUDA Kernel overall time: " << time_ms << " ms" << std::endl;

  cudaMemcpy(y_host, y, sizeof(float) * N, cudaMemcpyDeviceToHost);
  for(int i = 0; i < N; i++) {
    std::cout << "res " << y_host[i] << std::endl;
  }
  // Free memory
  cudaFree(x);
  cudaFree(y);

}

int main() {
    test2();
    std::cout << "end" << std::endl;
    return 0;
}

预期结果 #

每个核函数依次执行后，y数组的值应该逐步递增：

res 2.46
res 3.46
res 4.46
res 5.46
res 6.46
...

实际结果 #

所有结果都是最后一个核函数的计算值（1.23*2 +19=21.46）：

res 21.46
res 21.46
res 21.46
res 21.46
res 21.46
...

我尝试手动展开前几个核函数的参数，结果还是一样：

// Run graphs
for(int istep=0; istep<NSTEP; istep++){
  if(!graphCreated){
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
    for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
      if(ikrnl == 0)
        shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, 0);
      else if(ikrnl == 1)
        shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, 1);
      else if(ikrnl == 2)
        shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, 2);
      else
        shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, ikrnl);
    }
    cudaStreamEndCapture(stream, &graph);
    cudaGraphNode_t* nodes = NULL;
    size_t num_nodes = 0;
    CHECK(cudaGraphGetNodes(graph, nodes, &num_nodes));
    std::cout << "Num of nodes in the graph: " << num_nodes
              << std::endl;
    CHECK(cudaGraphInstantiate(&instance, graph, NULL, NULL, 0));
    graphCreated=true;
  }
  CHECK(cudaGraphLaunch(instance, stream));
  cudaStreamSynchronize(stream);
}  // End run graphs

问题原因 #

这是CUDA Graph的写后读依赖优化导致的结果。

CUDA Graph在捕获流中操作时，会自动分析数据依赖关系。你的代码中，所有核函数都对同一个设备数组y执行写操作，且没有中间的读操作或同步点。CUDA优化器会判定：前面的核函数对y的写操作会被后面的核函数完全覆盖，因此可以直接跳过前面所有核函数，只保留最后一个写操作的核函数。这就是为什么最终结果都是最后一个核函数的计算值。

手动展开参数后结果不变，是因为优化逻辑和参数无关，只和数据的读写依赖有关——只要所有核函数都是覆盖写入同一个数组，优化器就会合并这些操作。

解决方法 #

要让所有核函数都被执行，需要打破这种写覆盖的依赖关系，让CUDA Graph认为每个核函数的输出都是必要的。常见的方法有两种：

方法1：使用不同的输出数组 #

每个核函数写入不同的设备数组，避免覆盖：

// 预先分配多个输出数组
float *y_arr[NKERNEL];
for(int i=0; i<NKERNEL; i++){
    CHECK(cudaMalloc((void**)&y_arr[i], N*sizeof(float)));
}

// 捕获图时，每个核函数写入不同的数组
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
    shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y_arr[ikrnl], x, ikrnl);
}
cudaStreamEndCapture(stream, &graph);

方法2：添加显式依赖关系 #

在每个核函数之间添加同步标记，或者手动为图节点添加依赖，强制CUDA按顺序执行所有核函数：

cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
cudaEvent_t prev_event;
CHECK(cudaEventCreate(&prev_event));
for(int ikrnl=0; ikrnl<NKERNEL; ikrnl++){
    shortKernel<<<BLOCKS, THREADS, 0, stream>>>(y, x, ikrnl);
    cudaEventRecord(prev_event, stream);
    // 让下一个核函数等待当前核函数完成
    cudaStreamWaitEvent(stream, prev_event, 0);
}
cudaStreamEndCapture(stream, &graph);
cudaEventDestroy(prev_event);

或者在捕获图后手动添加节点依赖：

cudaStreamEndCapture(stream, &graph);
// 获取所有节点
size_t num_nodes;
CHECK(cudaGraphGetNodes(graph, NULL, &num_nodes));
cudaGraphNode_t* nodes = new cudaGraphNode_t[num_nodes];
CHECK(cudaGraphGetNodes(graph, nodes, &num_nodes));
// 依次添加依赖，让每个节点依赖前一个节点
for(int i=1; i<num_nodes; i++){
    CHECK(cudaGraphAddDependencies(graph, &nodes[i-1], &nodes[i], 1));
}
// 实例化图
CHECK(cudaGraphInstantiate(&instance, graph, NULL, NULL, 0));
delete[] nodes;

这样CUDA Graph就会按顺序执行所有核函数，不会优化掉前面的操作。

内容的提问来源于stack exchange，提问作者Jackson