CUDA Graphs: снижение оверхеда запуска при повторяющихся вычислениях

Введение

Когда приложение многократно выполняет одну и ту же последовательность операций CUDA (например, серию ядер в цикле обучения нейросети или итеративном численном алгоритме), накладные расходы CPU на запуск каждой отдельной операции (kernel launch overhead) могут стать заметным фактором, особенно если сами ядра выполняются быстро. CUDA Graphs позволяет один раз «записать» последовательность операций как граф зависимостей, а затем многократно запускать весь граф одним вызовом, существенно снижая накладные расходы на повторных итерациях. В статье разберём базовое использование CUDA Graphs.

Концепция

CUDA Graph — представление последовательности операций (запуски ядер, копирования памяти) как направленного ациклического графа с явными зависимостями между узлами. Граф можно построить двумя способами: явно через API построения графа (cudaGraphAddKernelNode и аналоги) или, что значительно проще на практике, через захват потока (stream capture): операции выполняются в обычном CUDA-потоке между вызовами cudaStreamBeginCapture() и cudaStreamEndCapture(), а CUDA автоматически записывает граф зависимостей вместо реального выполнения операций. После построения граф инстанцируется в исполняемую форму (cudaGraphInstantiate) и запускается многократно (cudaGraphLaunch) с минимальными накладными расходами CPU на каждый повторный запуск.

Пример кода

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

__global__ void stepA(float *data, int n)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) data[idx] += 1.0f;
}

__global__ void stepB(float *data, int n)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) data[idx] *= 2.0f;
}

int main()
{
    const int n = 1 << 20;
    float *deviceData;
    cudaMalloc(&deviceData, n * sizeof(float));
    cudaMemset(deviceData, 0, n * sizeof(float));

    cudaStream_t stream;
    cudaStreamCreate(&stream);

    const int threadsPerBlock = 256;
    const int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;

    cudaGraph_t graph;
    cudaGraphExec_t graphExec;

    // Захват последовательности операций в граф вместо реального выполнения
    cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
    stepA<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(deviceData, n);
    stepB<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(deviceData, n);
    cudaStreamEndCapture(stream, &graph);

    // Граф инстанцируется один раз...
    cudaGraphInstantiate(&graphExec, graph, nullptr, nullptr, 0);

    // ...и запускается многократно с минимальным оверхедом на запуск
    for (int iteration = 0; iteration < 1000; ++iteration) {
        cudaGraphLaunch(graphExec, stream);
    }
    cudaStreamSynchronize(stream);

    std::cout << "1000 итераций графа выполнены" << std::endl;

    cudaGraphExecDestroy(graphExec);
    cudaGraphDestroy(graph);
    cudaStreamDestroy(stream);
    cudaFree(deviceData);
    return 0;
}

Пояснения к коду

Между вызовами cudaStreamBeginCapture и cudaStreamEndCapture операции stepA и stepB, запущенные в потоке stream, не выполняются немедленно — вместо этого CUDA фиксирует их как узлы графа с зависимостью друг от друга (поскольку оба запущены в одном потоке последовательно). После cudaGraphInstantiate граф готов к многократному выполнению.

Цикл с 1000 итерациями cudaGraphLaunch запускает весь зафиксированный граф операций одним вызовом на каждой итерации — вместо 2000 отдельных запусков ядер (по два на итерацию) с соответствующими накладными расходами CPU на каждый, CPU выполняет всего 1000 легковесных вызовов запуска графа.

Подводные камни

• CUDA Graphs выгодны в первую очередь при множестве быстрых операций, где kernel launch overhead (порядка нескольких микросекунд на запуск) сравним или превышает время самого вычисления — для немногочисленных тяжёлых ядер выигрыш от графов незначителен и не оправдывает дополнительную сложность кода.
• Граф фиксирует не только последовательность операций, но и их параметры (адреса буферов, размеры) — если между запусками графа нужно использовать разные адреса памяти, нужно либо использовать cudaGraphExecKernelNodeSetParams для обновления параметров существующего графа, либо перестраивать граф заново, что снижает выгоду от подхода.
• Сложность отладки. Поскольку операции внутри stream capture не выполняются немедленно, традиционные инструменты отладки и профилирования сложнее применять напрямую к процессу захвата — отладку логики ядер обычно проводят отдельно, до перехода на граф-ориентированный запуск.
• Не все операции CUDA API поддерживают захват в граф. Некоторые операции (особенно связанные с управлением памятью динамически во время выполнения) не могут быть частью захватываемого потока — попытка их использования во время cudaStreamBeginCapture приведёт к ошибке захвата.