构造具有动态参数的CUDA图表

自从CUDA Graphs在CUDA 10中引入以来，CUDA Graph已经用于各种应用中。图形将一组CUDA内核和其他CUDA操作组合在一起，并使用指定的依赖关系树执行它们。它通过结合与CUDA内核启动和CUDA API调用相关的驱动程序活动来加快工作流。它还通过硬件加速强制实施依赖关系，而不是在可能的情况下仅依赖CUDA流和事件。

构造CUDA图有两种主要方法：显式API调用和流捕获。

使用显式API调用构造CUDA图

通过这种构建CUDA图的方法，由CUDA内核和CUDA内存操作形成的图节点通过调用cudaGraphAdd*节点API添加到图中，其中*被替换为节点类型。节点之间的依赖关系是用API显式设置的。

使用显式API构建CUDA图的好处是，cudaGraphAdd*Node API返回节点句柄（cudaGraph Node_t），可以用作未来节点更新的引用。例如，可以使用cudaGraphExecKernelNodeSetParams以最低成本更新实例化图中内核节点的内核启动配置和内核函数参数。

缺点是，在使用CUDA图加速现有代码的场景中，使用显式API调用构造CUDA图通常需要大量代码更改，尤其是有关代码的控制流和函数调用结构的更改。

使用流捕获构建CUDA图

通过这种构建CUDA图的方法，cudaStreamBeginCapture和cudaStream EndCapture被放置在代码块的前后。代码块启动的所有设备活动都会被记录、捕获并分组到CUDA图中。节点之间的依赖关系是从流捕获区域内的CUDA流或事件API调用中推断出来的。

使用流捕获构建CUDA图的好处是，对于现有代码，需要的代码更改更少。原始代码结构可以基本保持不变，图形构造是以自动方式执行的。

这种构建CUDA图的方法也有缺点。在流捕获区域内，所有内核启动配置和内核函数参数以及CUDA API调用参数都按值记录。每当任何配置和参数发生更改时，捕获的然后实例化的图形就会过期。

在《在动态环境中使用CUDA图》一文中提供了两种解决方案：

重新捕获工作流。当重新捕获的图与实例化的图具有相同的节点拓扑时，不需要重新实例化，并且可以使用cudaGraphExecUpdate执行整个图更新。

以配置和参数集作为键缓存CUDA图。每组配置和参数都与缓存中的不同CUDA图相关联。在运行工作流时，配置和参数集首先被抽象为一个键。然后在缓存中找到相应的图（如果它已经存在）并启动。

然而，在某些工作流中，两种解决方案都不能很好地工作。重新捕获然后更新方法在纸面上很有效，但在某些情况下，重新捕获和更新本身的成本很高。也有一些情况下，无法将每组参数与CUDA图相关联。例如，具有浮点数字参数的情况很难缓存，因为可能存在大量的浮点数字。

用显式API构造的CUDA图很容易更新，但这种方法可能过于繁琐，灵活性较差。CUDA图可以通过流捕获灵活地构造，但生成的图很难更新，而且更新成本很高。

综合方法

在本文中，我提供了一种使用显式API和流捕获方法构建CUDA图的方法，从而实现两者的优点，避免两者的缺点。

例如，在顺序启动三个内核的工作流中，前两个内核具有静态启动配置和参数，而最后一个内核具有动态启动配置和属性。

使用流捕获来记录前两个内核的启动，并调用显式API将最后一个内核节点添加到捕获图中。然后，显式API返回的节点句柄用于在每次启动图之前用动态配置和参数更新实例化图。

下面的代码示例说明了这个想法：

cudaStream_t stream;
std::vector _node_list;
cudaGraphExec_t _graph_exec;

if (not using_graph) {
  
  first_static_kernel<<<1, 1, 0, stream>>>(static_parameters);
  second_static_kernel<<<1, 1, 0, stream>>>(static_parameters);
  dynamic_kernel<<<1, 1, 0, stream>>>(dynamic_parameters);

} else {

  if (capturing_graph) {

    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

    first_static_kernel<<<1, 1, 0, stream>>>(static_parameters);
    second_static_kernel<<<1, 1, 0, stream>>>(static_parameters);

    // Get the current stream capturing graph
    cudaGraph_t _capturing_graph;
    cudaStreamCaptureStatus _capture_status;
    const cudaGraphNode_t *_deps;
    size_t _dep_count;
    cudaStreamGetCaptureInfo_v2(stream, &_capture_status, nullptr &_capturing_graph, &_deps, &_dep_count);

    // Manually add a new kernel node
    cudaGraphNode_t new_node;
    cudakernelNodeParams _dynamic_params_cuda;
    cudaGraphAddKernelNode(&new_node, _capturing_graph, _deps, _dep_count, &_dynamic_params_cuda);
    // ... and store the new node for future references
    _node_list.push_back(new_node);

    // Update the stream dependencies
    cudaStreamUpdateCaptureDependencies(stream, &new_node, 1, 1); 

    // End the capture and instantiate the graph
    cudaGraph_t _captured_graph;
    cudaStreamEndCapture(stream, &_captured_graph);

    cudaGraphInstantiate(&_graph_exec, _captured_graph, nullptr, nullptr, 0);

  } else if (updating_graph) {
    cudakernelNodeParams _dynamic_params_updated_cuda;
  
    cudaGraphExecKernelNodeSetParams(_graph_exec, _node_list[0], &_dynamic_params_updated_cuda);

  }
}

在此示例中，cudaStreamGetCaptureInfo_v2提取当前正在记录并捕获到的CUDA图形。在调用cudaStreamUpdateCaptureDependencies更新当前捕获流的依赖项树之前，会将一个内核节点添加到此图中，并返回和存储节点句柄（new_node）。最后一步是必要的，以确保随后捕获的任何其他活动都在这些手动添加的节点上正确设置了它们的依赖项。

使用这种方法，即使参数是动态的，也可以通过轻量级的cudaGraphExecKernelNodeSetParams调用直接重用相同的实例化图（cudaGraph Exec_t对象）。本文中的第一张图片显示了这种用法。

此外，捕获和更新代码路径可以组合成一段代码，与启动最后两个内核的原始代码相邻。这会造成最少的代码更改，并且不会破坏原始的控制流和函数调用结构。

新方法在带有动态参数的蜂树/cuda图独立代码示例中详细显示。cudaStreamGetCaptureInfo_v2和cudaStream UpdateCaptureDependencies是CUDA 11.3中引入的新CUDA运行时API。

绩效结果

使用带有动态参数的蜂巢树/cuda图独立代码示例，我用三种不同的方法测量了运行受内核启动开销约束的相同动态工作流的性能：

在没有CUDA图形加速的情况下运行

使用重新捕获然后更新方法运行CUDA图

使用本文介绍的组合方法运行CUDA图

表1显示了结果。本文中提到的方法的提速很大程度上取决于底层工作流。

结论

在本文中，我介绍了一种结合显式API和流捕获方法构建CUDA图的方法。它提供了一种以最低成本为具有动态参数的工作流重用实例化图的方法。

关于作者

Tu Jiqun在加入NVIDIA担任高级计算机开发技术工程师之前，曾获得哥伦比亚大学晶格QCD物理学博士学位。在NVIDIA，他致力于在最新的NVIDIAGPU上使用最新的硬件和软件功能，以加速广泛的HPC应用程序。

审核编辑：郭婷