具有条件节点的 CUDA 图的动态控制流

具有条件节点的 CUDA 图的动态控制流

引用

nvidia

1.

https://developer.nvidia.com/zh-cn/blog/dynamic-control-flow-in-cuda-graphs-with-conditional-nodes/

CUDA 图可以显著提高性能，因为驱动程序能够使用任务和依赖关系的完整描述来优化执行。特别是在静态工作流中，图形可以提供难以置信的好处，其创建的开销可以在多次连续启动中分摊，从而提高整体性能。然而，几乎所有的问题都涉及某种形式的决策，这可能需要分解图并将控制权返回给 CPU，以决定下一步要启动哪些工作。像这样分解工作会损害 CUDA 进行优化的能力，占用 CPU 资源，并增加每次图形启动的开销。从 CUDA 12.4 开始，CUDA Graphs 支持条件节点，这使得图形的部分能够有条件地或重复地执行，而不需要将控制返回到 CPU。这释放了 CPU 资源，使得更多的工作流能够在单个图形中表示，从而提高了计算效率。

条件节点

条件节点有两种风格：

• IF 节点：如果条件值为 true，则每次评估节点时执行一次主体。

• WHILE 节点：只要条件值为 true，当对节点求值时，就会重复执行主体。

条件节点是容器节点，类似于子图节点，但节点中包含的图的执行取决于条件变量的值。与节点关联的条件值由必须在节点之前创建的句柄访问，该句柄还可以指定在图形的每个开始处应用的初始化。可以通过调用在 CUDA 内核中设置条件值的

cudaGraphSetConditional

函数。

创建条件节点时，还会创建一个空图，并将句柄返回给用户。此图与节点绑定，并将根据条件值执行。此条件体图可以使用图形 API或者通过使用捕获异步 CUDA 调用的

cudaStreamBeginCaptureToGraph

函数。

条件节点也可以嵌套。例如，可以使用包含条件 IF 节点的体图创建条件 WHILE 节点。

条件节点体图可以包含以下任意一项：

• 内核节点（CNP，当前不支持协同）

• 空节点

• 子图节点

• Memset 节点

• Memcopy 节点

• 条件节点

这递归地应用于子图和条件体。所有内核，包括嵌套条件句或任何级别的子图中的内核，都必须属于同一个 CUDA 上下文。Memcopies 和 memset 必须作用于可从条件节点的上下文访问的内存。

完整的样品代码可在CUDA 样本库中找到。下一节将通过一些示例来展示如何处理条件节点的方法。

条件 IF 节点

如果条件为非零，则无论何时评估 IF 节点，都将执行一次 IF 节点的体图。图 1 描述了一个图，其中中间节点 B 是包含四节点图的 IF 条件节点：

图 1。条件 IF 节点

为了显示如何创建此图，以下示例使用条件节点 B 上游的内核节点 A，根据该内核所做的工作结果设置条件的值。条件的主体是使用图 API 填充的。

首先，定义节点 A 内核。这个内核根据用户执行的一些任意计算的结果来设置条件句柄。

__global__ void setHandle(cudaGraphConditionalHandle handle)  
{  
 unsigned int value = 0;  
 // We could perform some work here and set value based on the result of that work.  
 if (someCondition) {  
 // Set ‘value’ to non-zero if we want the conditional body to execute  
 value = 1;  
 }  
 cudaGraphSetConditional(handle, value);  
}  

接下来，定义一个函数来构建图。此函数分配条件句柄、创建节点并填充条件图的主体。为了清楚起见，省略了启动和执行图形的代码。

cudaGraph_t createGraph() {  
 cudaGraph_t graph;  
 cudaGraphNode_t node;  
 void *kernelArgs[1];  
 cudaGraphCreate(&graph, 0);  
 cudaGraphConditionalHandle handle;  
 cudaGraphConditionalHandleCreate(&handle, graph);  
 // Use a kernel upstream of the conditional to set the handle value  
 cudaGraphNodeParams kParams = { cudaGraphNodeTypeKernel };  
 kParams.kernel.func = (void *)setHandle;  
 kParams.kernel.gridDim.x = kParams.kernel.gridDim.y = kParams.kernel.gridDim.z = 1;  
 kParams.kernel.blockDim.x = kParams.kernel.blockDim.y = kParams.kernel.blockDim.z = 1;  
 kParams.kernel.kernelParams = kernelArgs;  
 kernelArgs[0] = &handle  
 cudaGraphAddNode(&node, graph, NULL, 0, &kParams);  
 cudaGraphNodeParams cParams = { cudaGraphNodeTypeConditional };  
 cParams.conditional.handle = handle;  
 cParams.conditional.type = cudaGraphCondTypeIf;  
 cParams.conditional.size = 1;  
 cudaGraphAddNode(&node, graph, &node, 1, &cParams);  
 cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];  
 // Populate the body of the conditional node  
 cudaGraphNode_t bodyNodes[4];  
 cudaGraphNodeParams params[4] = { ... }; // Setup kernel parameters as needed.  
 cudaGraphAddNode(&bodyNodes[0], bodyGraph, NULL, 0, &params[0]);  
 cudaGraphAddNode(&bodyNodes[1], bodyGraph, &bodyNodes[0], 1, &params[1]);  
 cudaGraphAddNode(&bodyNodes[2], bodyGraph, &bodyNodes[0], 1, &params[2]);  
 cudaGraphAddNode(&bodyNodes[3], bodyGraph, &bodyNodes[1], 2, &params[3]);  
 return graph;  
}  

条件 WHILE 节点

只要条件为非零，WHILE 节点的体图就会重复执行。将在执行节点时以及每次完成体图后评估条件。下图描述了一个三节点图，其中中间节点 B 是包含三节点图的 WHILE 条件节点。

图 2:条件 WHILE 节点

为了了解如何创建此图，以下示例将句柄的默认值设置为非零值，以便在默认情况下执行 WHILE 循环。将默认值设置为非零，并将条件值保留在条件上游的内核中不修改，这有效地产生了一个 do-while 循环，其中条件体总是至少执行一次。创建 WHILE 循环，其中循环体仅在条件为 true 时执行，需要执行一些计算并在节点 a 中适当设置条件句柄。

在上一个示例中，条件主体由图 API 填充。在本例中，使用流捕获填充条件的主体。

第一步是定义一个内核，在每次执行条件体时设置条件值。在本例中，句柄是基于递减计数器的值设置的。

__global__ void loopKernel(cudaGraphConditionalHandle handle)  
{  
 static int count = 10;  
 cudaGraphSetConditional(handle, --count ? 1 : 0);  
}  

接下来，定义一个函数来构建图。此函数分配条件句柄、创建节点并填充条件图的主体。为了清楚起见，省略了启动和执行图形的代码。

cudaGraph_t createGraph() {  
 cudaGraph_t graph;  
 cudaGraphNode_t nodes[3];  
 cudaGraphCreate(&graph, 0);  
 // Insert kernel node A  
 cudaGraphNodeParams params = ...;  
 cudaGraphAddNode(&nodes[0], graph, NULL, 0, &params);  
 cudaGraphConditionalHandle handle;  
 cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault);  
 // Insert conditional node B  
 cudaGraphNodeParams cParams = { cudaGraphNodeTypeConditional };  
 cParams.conditional.handle = handle;  
 cParams.conditional.type = cudaGraphCondTypeWhile;  
 cParams.conditional.size = 1;  
 cudaGraphAddNode(&nodes[1], graph, &nodes[0], 1, &cParams);  
 cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];  
 cudaStream_t captureStream;  
 cudaStreamCreate(&captureStream);  
 // Fill out body graph with stream capture.  
 cudaStreamBeginCaptureToGraph(captureStream,  
 bodyGraph,  
 nullptr,  
 nullptr,  
 0,  
 cudaStreamCaptureModeRelaxed);  
 myKernel1<<<..., captureStream>>>(...);  
 myKernel2<<<..., captureStream>>>(...);  
 loopKernel<<<1, 1, 0, captureStream>>>(handle);  
 cudaStreamEndCapture(captureStream, nullptr);  
 cudaStreamDestroy(captureStream);  
 // Insert kernel node C.  
 params = ...;  
 cudaGraphAddNode(&nodes[2], graph, &nodes[1], 1, &params);  
 return graph;  
}  

此示例使用

cudaStreamBeginCaptureToGraph

，这是 CUDA 12.3 中新增加的一个 API，它使流捕获能够将节点插入到现有的图中。借助这个 API，可以将多个单独的捕获组合到一个图形对象中。此外，这个 API 还允许填充与条件节点一起创建的条件体图对象。

结论

CUDA Graphs 为静态工作流提供了难以置信的好处，在静态工作流中，图形创建的开销可以在多次连续启动中分摊。消除对图形的分割并将控制权返回给 CPU，以决定优先启动哪个，这有助于减少 CPU 开销和延迟。使用具有条件节点的 CUDA 图形可以有条件地或重复地执行图形的部分，而无需将控制权返回给 CPU。这释放了 CPU 资源，并使单个图形能够表示更复杂的工作流。

欲了解更多关于条件节点的信息，请参阅CUDA 编程指南。要探索简单、完整的示例代码，请访问NVIDIA/cuda-samples在 GitHub 上。同时，您也可以加入NVIDIA 开发者 CUDA 论坛中的讨论。

本文原文来自NVIDIA开发者博客