Graph kernel nodes don't keep kernel argument objects alive #2039

def test_kernel_args_buffer_not_kept_alive():
    """Buffer passed as a kernel arg is not kept alive by the graph.

    The ParamHolder built in GN_launch holds the only references to the
    kernel args and is destroyed when GN_launch returns. With no other
    reference to the Buffer, its device memory is freed and the graph
    is left with a stale pointer.
    """
    import ctypes
    import gc
    import weakref

    from cuda.bindings import driver as drv
    from cuda.core import Device, DeviceMemoryResource, LaunchConfig, Program, ProgramOptions
    from cuda.core.graph import GraphDefinition

    dev = Device()
    dev.set_current()
    stream = dev.create_stream()
    mr = DeviceMemoryResource(dev)

    arch = "".join(f"{i}" for i in dev.compute_capability)
    prog = Program(
        "__global__ void add_one(int *a) { *a += 1; }",
        code_type="c++",
        options=ProgramOptions(std="c++17", arch=f"sm_{arch}"),
    )
    add_one = prog.compile("cubin", name_expressions=("add_one",)).get_kernel("add_one")
    config = LaunchConfig(grid=1, block=1)

    buf = mr.allocate(ctypes.sizeof(ctypes.c_int), stream=stream)
    stream.sync()
    buf_weak = weakref.ref(buf)
    drv.cuMemcpyHtoD(int(buf.handle), (ctypes.c_int * 1)(0), ctypes.sizeof(ctypes.c_int))

    g = GraphDefinition()
    g.launch(config, add_one, buf)

    del buf
    gc.collect()
    assert buf_weak() is None  # nothing in the graph kept the Buffer alive

    g.instantiate().launch(stream)
    stream.sync()  # crash or memory corruption