cuda.core: ending a GraphBuilder after an invalidated conditional-body capture crashes the process (CUDA driver bug)

// Reproducer: ending an invalidated capture that targets a conditional node's
// body graph destroys that body graph inside the driver, and tearing down the
// parent graph destroys it a second time.
//
// Sequence (all calls legal except the one marked ILLEGAL):
//   1. cuStreamBeginCapture on a parent stream.
//   2. Add a conditional IF node to the captured graph (cuGraphAddNode); the
//      driver allocates the node's body graph and owns it through the node.
//   3. cuStreamBeginCaptureToGraph into that body graph on a second stream.
//   4. ILLEGAL: cuCtxSynchronize during capture. Returns
//      CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED and invalidates every capture in
//      the context. Nothing is destroyed yet.
//   5. cuStreamEndCapture on the body stream. Returns
//      CUDA_ERROR_STREAM_CAPTURE_INVALIDATED and a NULL graph, as documented,
//      and (undocumented) destroys the body graph the conditional node owns.
//   6. cuStreamEndCapture on the parent stream. Returns the same error and
//      destroys the parent graph, whose conditional node destroys the body
//      graph again: double free.
// --order parent-first swaps 5 and 6: the parent's teardown frees the body
// graph while the body stream still captures into it, and step 5 then touches
// freed memory.
// --control skips step 4: both captures end cleanly and cuGraphDestroy on the
// parent frees everything exactly once.
//
// Build: gcc -O0 -g -Wall -o cond_body_double_free cond_body_double_free.c
//            -I$CUDA_HOME/include -L$CUDA_HOME/lib/stubs -lcuda
// Run:   ./cond_body_double_free [--order body-first|parent-first]
//                                [--mode global|thread_local|relaxed] [--control]
// A plain run may or may not abort; run under valgrind or glibc's malloc
// debugging (LD_PRELOAD=libc_malloc_debug.so.0 GLIBC_TUNABLES=glibc.malloc.check=3)
// to make the second free visible.

#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static const char* err_name(CUresult r) {
    const char* s = "?";
    cuGetErrorName(r, &s);
    return s;
}

#define MUST(call)                                                             \
    do {                                                                       \
        CUresult _r = (call);                                                  \
        if (_r != CUDA_SUCCESS) {                                              \
            fprintf(stderr, "line %d: %s -> %s\n", __LINE__, #call, err_name(_r)); \
            exit(2);                                                           \
        }                                                                      \
    } while (0)

#define SHOW(call)                                                             \
    do {                                                                       \
        CUresult _r = (call);                                                  \
        printf("%-64s -> %s\n", #call, err_name(_r));                          \
        fflush(stdout);                                                        \
    } while (0)

int main(int argc, char** argv) {
    int body_first = 1;
    int control = 0;
    CUstreamCaptureMode mode = CU_STREAM_CAPTURE_MODE_GLOBAL;
    const char* mode_name = "global";
    for (int i = 1; i < argc; ++i) {
        if (!strcmp(argv[i], "--order") && i + 1 < argc) {
            body_first = strcmp(argv[++i], "parent-first") != 0;
        } else if (!strcmp(argv[i], "--mode") && i + 1 < argc) {
            mode_name = argv[++i];
            if (!strcmp(mode_name, "thread_local")) {
                mode = CU_STREAM_CAPTURE_MODE_THREAD_LOCAL;
            } else if (!strcmp(mode_name, "relaxed")) {
                mode = CU_STREAM_CAPTURE_MODE_RELAXED;
            } else if (strcmp(mode_name, "global")) {
                fprintf(stderr, "unknown mode %s\n", mode_name);
                return 2;
            }
        } else if (!strcmp(argv[i], "--control")) {
            control = 1;
        } else {
            fprintf(stderr,
                    "usage: %s [--order body-first|parent-first] "
                    "[--mode global|thread_local|relaxed] [--control]\n",
                    argv[0]);
            return 2;
        }
    }

    int drv = 0;
    MUST(cuInit(0));
    MUST(cuDriverGetVersion(&drv));
    printf("driver %d.%d  order=%s  mode=%s  control=%d\n", drv / 1000, (drv % 1000) / 10,
           body_first ? "body-first" : "parent-first", mode_name, control);

    CUdevice dev;
    CUcontext ctx;
    MUST(cuDeviceGet(&dev, 0));
    MUST(cuDevicePrimaryCtxRetain(&ctx, dev));
    MUST(cuCtxSetCurrent(ctx));

    CUstream parent_stream, body_stream;
    MUST(cuStreamCreate(&parent_stream, CU_STREAM_NON_BLOCKING));
    MUST(cuStreamCreate(&body_stream, CU_STREAM_NON_BLOCKING));
    CUdeviceptr dptr;
    MUST(cuMemAlloc(&dptr, 4096));

    // 1. Parent capture.
    MUST(cuStreamBeginCapture(parent_stream, mode));
    CUstreamCaptureStatus st;
    CUgraph parent = NULL;
    const CUgraphNode* deps = NULL;
    size_t ndeps = 0;
    MUST(cuStreamGetCaptureInfo(parent_stream, &st, NULL, &parent, &deps, NULL, &ndeps));

    // 2. Conditional IF node in the captured graph. The driver allocates the
    //    body graph and hands it back in phGraph_out; the node owns it.
    CUgraphConditionalHandle handle;
    MUST(cuGraphConditionalHandleCreate(&handle, parent, ctx, 1, CU_GRAPH_COND_ASSIGN_DEFAULT));
    CUgraphNodeParams params;
    memset(&params, 0, sizeof params);
    params.type = CU_GRAPH_NODE_TYPE_CONDITIONAL;
    params.conditional.handle = handle;
    params.conditional.type = CU_GRAPH_COND_TYPE_IF;
    params.conditional.size = 1;
    params.conditional.ctx = ctx;
    CUgraphNode cond_node;
    MUST(cuGraphAddNode(&cond_node, parent, deps, NULL, ndeps, &params));
    MUST(cuStreamUpdateCaptureDependencies(parent_stream, &cond_node, NULL, 1,
                                           CU_STREAM_SET_CAPTURE_DEPENDENCIES));
    CUgraph body = params.conditional.phGraph_out[0];
    printf("parent graph %p, conditional node %p, body graph %p\n", (void*)parent,
           (void*)cond_node, (void*)body);

    // 3. Capture into the body graph on its own stream, with one node in it.
    MUST(cuStreamBeginCaptureToGraph(body_stream, body, NULL, NULL, 0, mode));
    MUST(cuMemsetD8Async(dptr, 0, 4096, body_stream));

    CUgraph out;
    if (control) {
        out = NULL;
        SHOW(cuStreamEndCapture(body_stream, &out));
        printf("  body out=%p (is the body graph: %d)\n", (void*)out, out == body);
        out = NULL;
        SHOW(cuStreamEndCapture(parent_stream, &out));
        printf("  parent out=%p (is the parent graph: %d)\n", (void*)out, out == parent);
        SHOW(cuGraphDestroy(parent));
    } else {
        // 4. ILLEGAL during capture: invalidates every capture in the context.
        SHOW(cuCtxSynchronize());
        SHOW(cuStreamIsCapturing(body_stream, &st));
        printf("  body stream capture status = %d (2 = invalidated, not ended)\n", (int)st);
        SHOW(cuStreamIsCapturing(parent_stream, &st));
        printf("  parent stream capture status = %d\n", (int)st);

        // 5 and 6. End both captures; each is the documented cleanup call.
        if (body_first) {
            out = (CUgraph)1;
            SHOW(cuStreamEndCapture(body_stream, &out));
            printf("  out=%p\n", (void*)out);
            out = (CUgraph)1;
            SHOW(cuStreamEndCapture(parent_stream, &out));
            printf("  out=%p\n", (void*)out);
        } else {
            out = (CUgraph)1;
            SHOW(cuStreamEndCapture(parent_stream, &out));
            printf("  out=%p\n", (void*)out);
            out = (CUgraph)1;
            SHOW(cuStreamEndCapture(body_stream, &out));
            printf("  out=%p\n", (void*)out);
        }
    }

    SHOW(cuStreamDestroy(body_stream));
    SHOW(cuStreamDestroy(parent_stream));
    SHOW(cuMemFree(dptr));
    SHOW(cuDevicePrimaryCtxRelease(dev));
    printf("REACHED_END\n");
    return 0;
}