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 条件节点:
为了显示如何创建此图,以下示例使用条件节点 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, ¶ms[0]);
cudaGraphAddNode(&bodyNodes[1], bodyGraph, &bodyNodes[0], 1, ¶ms[1]);
cudaGraphAddNode(&bodyNodes[2], bodyGraph, &bodyNodes[0], 1, ¶ms[2]);
cudaGraphAddNode(&bodyNodes[3], bodyGraph, &bodyNodes[1], 2, ¶ms[3]);
return graph;
}
条件 WHILE 节点
只要条件为非零,WHILE 节点的体图就会重复执行。将在执行节点时以及每次完成体图后评估条件。下图描述了一个三节点图,其中中间节点 B 是包含三节点图的 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, ¶ms);
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, ¶ms);
return graph;
}
此示例使用 cudaStreamBeginCaptureToGraph
,这是 CUDA 12.3 中新增加的一个 API,它使流捕获能够将节点插入到现有的图中。借助这个 API,可以将多个单独的捕获组合到一个图形对象中。此外,这个 API 还允许填充与条件节点一起创建的条件体图对象。
结论
CUDA Graphs 为静态工作流提供了难以置信的好处,在静态工作流中,图形创建的开销可以在多次连续启动中分摊。消除对图形的分割并将控制权返回给 CPU,以决定优先启动哪个,这有助于减少 CPU 开销和延迟。使用具有条件节点的 CUDA 图形可以有条件地或重复地执行图形的部分,而无需将控制权返回给 CPU。这释放了 CPU 资源,并使单个图形能够表示更复杂的工作流。
欲了解更多关于条件节点的信息,请参阅 CUDA 编程指南。要探索简单、完整的示例代码,请访问 NVIDIA/cuda-samples 在 GitHub 上。同时,您也可以加入 NVIDIA 开发者 CUDA 论坛 中的讨论。