構(gòu)造具有動(dòng)態(tài)參數(shù)的CUDA圖表

自從CUDA Graphs在CUDA 10中引入以來(lái)，CUDA Graph已經(jīng)用于各種應(yīng)用中。圖形將一組CUDA內(nèi)核和其他CUDA操作組合在一起，并使用指定的依賴關(guān)系樹(shù)執(zhí)行它們。它通過(guò)結(jié)合與CUDA內(nèi)核啟動(dòng)和CUDA API調(diào)用相關(guān)的驅(qū)動(dòng)程序活動(dòng)來(lái)加快工作流。它還通過(guò)硬件加速?gòu)?qiáng)制實(shí)施依賴關(guān)系，而不是在可能的情況下僅依賴CUDA流和事件。

構(gòu)造CUDA圖有兩種主要方法：顯式API調(diào)用和流捕獲。

使用顯式API調(diào)用構(gòu)造CUDA圖

通過(guò)這種構(gòu)建CUDA圖的方法，由CUDA內(nèi)核和CUDA內(nèi)存操作形成的圖節(jié)點(diǎn)通過(guò)調(diào)用cudaGraphAdd*節(jié)點(diǎn)API添加到圖中，其中*被替換為節(jié)點(diǎn)類型。節(jié)點(diǎn)之間的依賴關(guān)系是用API顯式設(shè)置的。

使用顯式API構(gòu)建CUDA圖的好處是，cudaGraphAdd*Node API返回節(jié)點(diǎn)句柄（cudaGraph Node_t），可以用作未來(lái)節(jié)點(diǎn)更新的引用。例如，可以使用cudaGraphExecKernelNodeSetParams以最低成本更新實(shí)例化圖中內(nèi)核節(jié)點(diǎn)的內(nèi)核啟動(dòng)配置和內(nèi)核函數(shù)參數(shù)。

缺點(diǎn)是，在使用CUDA圖加速現(xiàn)有代碼的場(chǎng)景中，使用顯式API調(diào)用構(gòu)造CUDA圖通常需要大量代碼更改，尤其是有關(guān)代碼的控制流和函數(shù)調(diào)用結(jié)構(gòu)的更改。

使用流捕獲構(gòu)建CUDA圖

通過(guò)這種構(gòu)建CUDA圖的方法，cudaStreamBeginCapture和cudaStream EndCapture被放置在代碼塊的前后。代碼塊啟動(dòng)的所有設(shè)備活動(dòng)都會(huì)被記錄、捕獲并分組到CUDA圖中。節(jié)點(diǎn)之間的依賴關(guān)系是從流捕獲區(qū)域內(nèi)的CUDA流或事件API調(diào)用中推斷出來(lái)的。

使用流捕獲構(gòu)建CUDA圖的好處是，對(duì)于現(xiàn)有代碼，需要的代碼更改更少。原始代碼結(jié)構(gòu)可以基本保持不變，圖形構(gòu)造是以自動(dòng)方式執(zhí)行的。

這種構(gòu)建CUDA圖的方法也有缺點(diǎn)。在流捕獲區(qū)域內(nèi)，所有內(nèi)核啟動(dòng)配置和內(nèi)核函數(shù)參數(shù)以及CUDA API調(diào)用參數(shù)都按值記錄。每當(dāng)任何配置和參數(shù)發(fā)生更改時(shí)，捕獲的然后實(shí)例化的圖形就會(huì)過(guò)期。

在《在動(dòng)態(tài)環(huán)境中使用CUDA圖》一文中提供了兩種解決方案：

重新捕獲工作流。當(dāng)重新捕獲的圖與實(shí)例化的圖具有相同的節(jié)點(diǎn)拓?fù)鋾r(shí)，不需要重新實(shí)例化，并且可以使用cudaGraphExecUpdate執(zhí)行整個(gè)圖更新。

以配置和參數(shù)集作為鍵緩存CUDA圖。每組配置和參數(shù)都與緩存中的不同CUDA圖相關(guān)聯(lián)。在運(yùn)行工作流時(shí)，配置和參數(shù)集首先被抽象為一個(gè)鍵。然后在緩存中找到相應(yīng)的圖（如果它已經(jīng)存在）并啟動(dòng)。

然而，在某些工作流中，兩種解決方案都不能很好地工作。重新捕獲然后更新方法在紙面上很有效，但在某些情況下，重新捕獲和更新本身的成本很高。也有一些情況下，無(wú)法將每組參數(shù)與CUDA圖相關(guān)聯(lián)。例如，具有浮點(diǎn)數(shù)字參數(shù)的情況很難緩存，因?yàn)榭赡艽嬖诖罅康母↑c(diǎn)數(shù)字。

用顯式API構(gòu)造的CUDA圖很容易更新，但這種方法可能過(guò)于繁瑣，靈活性較差。CUDA圖可以通過(guò)流捕獲靈活地構(gòu)造，但生成的圖很難更新，而且更新成本很高。

綜合方法

在本文中，我提供了一種使用顯式API和流捕獲方法構(gòu)建CUDA圖的方法，從而實(shí)現(xiàn)兩者的優(yōu)點(diǎn)，避免兩者的缺點(diǎn)。

例如，在順序啟動(dòng)三個(gè)內(nèi)核的工作流中，前兩個(gè)內(nèi)核具有靜態(tài)啟動(dòng)配置和參數(shù)，而最后一個(gè)內(nèi)核具有動(dòng)態(tài)啟動(dòng)配置和屬性。

使用流捕獲來(lái)記錄前兩個(gè)內(nèi)核的啟動(dòng)，并調(diào)用顯式API將最后一個(gè)內(nèi)核節(jié)點(diǎn)添加到捕獲圖中。然后，顯式API返回的節(jié)點(diǎn)句柄用于在每次啟動(dòng)圖之前用動(dòng)態(tài)配置和參數(shù)更新實(shí)例化圖。

下面的代碼示例說(shuō)明了這個(gè)想法：

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提取當(dāng)前正在記錄并捕獲到的CUDA圖形。在調(diào)用cudaStreamUpdateCaptureDependencies更新當(dāng)前捕獲流的依賴項(xiàng)樹(shù)之前，會(huì)將一個(gè)內(nèi)核節(jié)點(diǎn)添加到此圖中，并返回和存儲(chǔ)節(jié)點(diǎn)句柄（new_node）。最后一步是必要的，以確保隨后捕獲的任何其他活動(dòng)都在這些手動(dòng)添加的節(jié)點(diǎn)上正確設(shè)置了它們的依賴項(xiàng)。

使用這種方法，即使參數(shù)是動(dòng)態(tài)的，也可以通過(guò)輕量級(jí)的cudaGraphExecKernelNodeSetParams調(diào)用直接重用相同的實(shí)例化圖（cudaGraph Exec_t對(duì)象）。本文中的第一張圖片顯示了這種用法。

此外，捕獲和更新代碼路徑可以組合成一段代碼，與啟動(dòng)最后兩個(gè)內(nèi)核的原始代碼相鄰。這會(huì)造成最少的代碼更改，并且不會(huì)破壞原始的控制流和函數(shù)調(diào)用結(jié)構(gòu)。

新方法在帶有動(dòng)態(tài)參數(shù)的蜂樹(shù)/cuda圖獨(dú)立代碼示例中詳細(xì)顯示。cudaStreamGetCaptureInfo_v2和cudaStream UpdateCaptureDependencies是CUDA 11.3中引入的新CUDA運(yùn)行時(shí)API。

績(jī)效結(jié)果

使用帶有動(dòng)態(tài)參數(shù)的蜂巢樹(shù)/cuda圖獨(dú)立代碼示例，我用三種不同的方法測(cè)量了運(yùn)行受內(nèi)核啟動(dòng)開(kāi)銷約束的相同動(dòng)態(tài)工作流的性能：

在沒(méi)有CUDA圖形加速的情況下運(yùn)行

使用重新捕獲然后更新方法運(yùn)行CUDA圖

使用本文介紹的組合方法運(yùn)行CUDA圖

表1顯示了結(jié)果。本文中提到的方法的提速很大程度上取決于底層工作流。

結(jié)論

在本文中，我介紹了一種結(jié)合顯式API和流捕獲方法構(gòu)建CUDA圖的方法。它提供了一種以最低成本為具有動(dòng)態(tài)參數(shù)的工作流重用實(shí)例化圖的方法。

關(guān)于作者

Tu Jiqun在加入NVIDIA擔(dān)任高級(jí)計(jì)算機(jī)開(kāi)發(fā)技術(shù)工程師之前，曾獲得哥倫比亞大學(xué)晶格QCD物理學(xué)博士學(xué)位。在NVIDIA，他致力于在最新的NVIDIAGPU上使用最新的硬件和軟件功能，以加速?gòu)V泛的HPC應(yīng)用程序。

審核編輯：郭婷