feat: add support for CUcheckpointRestoreArgs (NVIDIA#975)

* feat: add support for CUcheckpointRestoreArgs

* feat: generate APIs that use the newly allowed struct

* chore: replace to avoid failure to parse

* chore: fmt

* Update release notes for CUDA bindings 13.X.Y

---------

Co-authored-by: Leo Fang <leof@nvidia.com>

Author: cpcloud

cuda_bindings/cuda/bindings/_bindings/cydriver.pxd.in

@@ -2239,6 +2239,11 @@ cdef CUresult _cuCheckpointProcessLock(int pid, CUcheckpointLockArgs* args) exce
 cdef CUresult _cuCheckpointProcessCheckpoint(int pid, CUcheckpointCheckpointArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil
 {{endif}}
 
+{{if 'cuCheckpointProcessRestore' in found_functions}}
+
+cdef CUresult _cuCheckpointProcessRestore(int pid, CUcheckpointRestoreArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil
+{{endif}}
+
 {{if 'cuCheckpointProcessUnlock' in found_functions}}
 
 cdef CUresult _cuCheckpointProcessUnlock(int pid, CUcheckpointUnlockArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil

cuda_bindings/cuda/bindings/_bindings/cydriver.pyx.in

@@ -463,6 +463,7 @@ cdef bint __cuPythonInit = False
 {{if 'cuCheckpointProcessGetState' in found_functions}}cdef void *__cuCheckpointProcessGetState = NULL{{endif}}
 {{if 'cuCheckpointProcessLock' in found_functions}}cdef void *__cuCheckpointProcessLock = NULL{{endif}}
 {{if 'cuCheckpointProcessCheckpoint' in found_functions}}cdef void *__cuCheckpointProcessCheckpoint = NULL{{endif}}
+{{if 'cuCheckpointProcessRestore' in found_functions}}cdef void *__cuCheckpointProcessRestore = NULL{{endif}}
 {{if 'cuCheckpointProcessUnlock' in found_functions}}cdef void *__cuCheckpointProcessUnlock = NULL{{endif}}
 {{if 'cuProfilerStart' in found_functions}}cdef void *__cuProfilerStart = NULL{{endif}}
 {{if 'cuProfilerStop' in found_functions}}cdef void *__cuProfilerStop = NULL{{endif}}
@@ -2667,6 +2668,10 @@ cdef int _cuPythonInit() except -1 nogil:
             global __cuCheckpointProcessCheckpoint
             _F_cuGetProcAddress_v2('cuCheckpointProcessCheckpoint', &__cuCheckpointProcessCheckpoint, 12080, CU_GET_PROC_ADDRESS_DEFAULT, NULL)
             {{endif}}
+            {{if 'cuCheckpointProcessRestore' in found_functions}}
+            global __cuCheckpointProcessRestore
+            _F_cuGetProcAddress_v2('cuCheckpointProcessRestore', &__cuCheckpointProcessRestore, 12080, CU_GET_PROC_ADDRESS_DEFAULT, NULL)
+            {{endif}}
             {{if 'cuCheckpointProcessUnlock' in found_functions}}
             global __cuCheckpointProcessUnlock
             _F_cuGetProcAddress_v2('cuCheckpointProcessUnlock', &__cuCheckpointProcessUnlock, 12080, CU_GET_PROC_ADDRESS_DEFAULT, NULL)
@@ -4895,6 +4900,10 @@ cdef int _cuPythonInit() except -1 nogil:
         global __cuCheckpointProcessCheckpoint
         __cuCheckpointProcessCheckpoint = windll.GetProcAddress(handle, 'cuCheckpointProcessCheckpoint')
         {{endif}}
+        {{if 'cuCheckpointProcessRestore' in found_functions}}
+        global __cuCheckpointProcessRestore
+        __cuCheckpointProcessRestore = windll.GetProcAddress(handle, 'cuCheckpointProcessRestore')
+        {{endif}}
         {{if 'cuCheckpointProcessUnlock' in found_functions}}
         global __cuCheckpointProcessUnlock
         __cuCheckpointProcessUnlock = windll.GetProcAddress(handle, 'cuCheckpointProcessUnlock')
@@ -7120,6 +7129,10 @@ cdef int _cuPythonInit() except -1 nogil:
         global __cuCheckpointProcessCheckpoint
         __cuCheckpointProcessCheckpoint = dlfcn.dlsym(handle, 'cuCheckpointProcessCheckpoint')
         {{endif}}
+        {{if 'cuCheckpointProcessRestore' in found_functions}}
+        global __cuCheckpointProcessRestore
+        __cuCheckpointProcessRestore = dlfcn.dlsym(handle, 'cuCheckpointProcessRestore')
+        {{endif}}
         {{if 'cuCheckpointProcessUnlock' in found_functions}}
         global __cuCheckpointProcessUnlock
         __cuCheckpointProcessUnlock = dlfcn.dlsym(handle, 'cuCheckpointProcessUnlock')
@@ -12583,6 +12596,18 @@ cdef CUresult _cuCheckpointProcessCheckpoint(int pid, CUcheckpointCheckpointArgs
     return err
 {{endif}}
 
+{{if 'cuCheckpointProcessRestore' in found_functions}}
+
+cdef CUresult _cuCheckpointProcessRestore(int pid, CUcheckpointRestoreArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil:
+    global __cuCheckpointProcessRestore
+    cuPythonInit()
+    if __cuCheckpointProcessRestore == NULL:
+        with gil:
+            raise RuntimeError('Function "cuCheckpointProcessRestore" not found')
+    err = (<CUresult (*)(int, CUcheckpointRestoreArgs*) except ?CUDA_ERROR_NOT_FOUND nogil> __cuCheckpointProcessRestore)(pid, args)
+    return err
+{{endif}}
+
 {{if 'cuCheckpointProcessUnlock' in found_functions}}
 
 cdef CUresult _cuCheckpointProcessUnlock(int pid, CUcheckpointUnlockArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil:
@@ -15986,6 +16011,13 @@ cpdef dict _inspect_function_pointers():
     data["__cuCheckpointProcessCheckpoint"] = <intptr_t>0
     {{endif}}
 
+    {{if 'cuCheckpointProcessRestore' in found_functions}}
+    global __cuCheckpointProcessRestore
+    data["__cuCheckpointProcessRestore"] = <intptr_t>__cuCheckpointProcessRestore
+    {{else}}
+    data["__cuCheckpointProcessRestore"] = <intptr_t>0
+    {{endif}}
+
     {{if 'cuCheckpointProcessUnlock' in found_functions}}
     global __cuCheckpointProcessUnlock
     data["__cuCheckpointProcessUnlock"] = <intptr_t>__cuCheckpointProcessUnlock

cuda_bindings/cuda/bindings/cydriver.pxd.in

@@ -2357,6 +2357,14 @@ cdef extern from "cuda.h":
 
     ctypedef CUcheckpointGpuPair_st CUcheckpointGpuPair
 
+    cdef struct CUcheckpointRestoreArgs_st:
+        CUcheckpointGpuPair* gpuPairs
+        unsigned int gpuPairsCount
+        char reserved[44]
+        cuuint64_t reserved1
+
+    ctypedef CUcheckpointRestoreArgs_st CUcheckpointRestoreArgs
+
     cdef struct CUcheckpointUnlockArgs_st:
         cuuint64_t reserved[8]
 
@@ -4907,6 +4915,11 @@ cdef CUresult cuCheckpointProcessLock(int pid, CUcheckpointLockArgs* args) excep
 cdef CUresult cuCheckpointProcessCheckpoint(int pid, CUcheckpointCheckpointArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil
 {{endif}}
 
+{{if 'cuCheckpointProcessRestore' in found_functions}}
+
+cdef CUresult cuCheckpointProcessRestore(int pid, CUcheckpointRestoreArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil
+{{endif}}
+
 {{if 'cuCheckpointProcessUnlock' in found_functions}}
 
 cdef CUresult cuCheckpointProcessUnlock(int pid, CUcheckpointUnlockArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil

cuda_bindings/cuda/bindings/cydriver.pyx.in

@@ -2686,6 +2686,12 @@ cdef CUresult cuCheckpointProcessCheckpoint(int pid, CUcheckpointCheckpointArgs*
     return cydriver._cuCheckpointProcessCheckpoint(pid, args)
 {{endif}}
 
+{{if 'cuCheckpointProcessRestore' in found_functions}}
+
+cdef CUresult cuCheckpointProcessRestore(int pid, CUcheckpointRestoreArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil:
+    return cydriver._cuCheckpointProcessRestore(pid, args)
+{{endif}}
+
 {{if 'cuCheckpointProcessUnlock' in found_functions}}
 
 cdef CUresult cuCheckpointProcessUnlock(int pid, CUcheckpointUnlockArgs* args) except ?CUDA_ERROR_NOT_FOUND nogil:

cuda_bindings/cuda/bindings/driver.pxd.in

@@ -5058,6 +5058,47 @@ cdef class CUcheckpointGpuPair_st:
     cdef CUuuid _newUuid
     {{endif}}
 {{endif}}
+{{if 'CUcheckpointRestoreArgs_st' in found_struct}}
+
+cdef class CUcheckpointRestoreArgs_st:
+    """
+    CUDA checkpoint optional restore arguments
+
+    Attributes
+    ----------
+    {{if 'CUcheckpointRestoreArgs_st.gpuPairs' in found_struct}}
+    gpuPairs : CUcheckpointGpuPair
+        Pointer to array of gpu pairs that indicate how to remap GPUs
+        during restore
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.gpuPairsCount' in found_struct}}
+    gpuPairsCount : unsigned int
+        Number of gpu pairs to remap
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
+    reserved : bytes
+        Reserved for future use, must be zeroed
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
+    reserved1 : cuuint64_t
+        Reserved for future use, must be zeroed
+    {{endif}}
+
+    Methods
+    -------
+    getPtr()
+        Get memory address of class instance
+    """
+    cdef cydriver.CUcheckpointRestoreArgs_st _pvt_val
+    cdef cydriver.CUcheckpointRestoreArgs_st* _pvt_ptr
+    {{if 'CUcheckpointRestoreArgs_st.gpuPairs' in found_struct}}
+    cdef size_t _gpuPairs_length
+    cdef cydriver.CUcheckpointGpuPair* _gpuPairs
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
+    cdef cuuint64_t _reserved1
+    {{endif}}
+{{endif}}
 {{if 'CUcheckpointUnlockArgs_st' in found_struct}}
 
 cdef class CUcheckpointUnlockArgs_st:
@@ -10563,6 +10604,39 @@ cdef class CUcheckpointGpuPair(CUcheckpointGpuPair_st):
     """
     pass
 {{endif}}
+{{if 'CUcheckpointRestoreArgs' in found_types}}
+
+cdef class CUcheckpointRestoreArgs(CUcheckpointRestoreArgs_st):
+    """
+    CUDA checkpoint optional restore arguments
+
+    Attributes
+    ----------
+    {{if 'CUcheckpointRestoreArgs_st.gpuPairs' in found_struct}}
+    gpuPairs : CUcheckpointGpuPair
+        Pointer to array of gpu pairs that indicate how to remap GPUs
+        during restore
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.gpuPairsCount' in found_struct}}
+    gpuPairsCount : unsigned int
+        Number of gpu pairs to remap
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
+    reserved : bytes
+        Reserved for future use, must be zeroed
+    {{endif}}
+    {{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
+    reserved1 : cuuint64_t
+        Reserved for future use, must be zeroed
+    {{endif}}
+
+    Methods
+    -------
+    getPtr()
+        Get memory address of class instance
+    """
+    pass
+{{endif}}
 {{if 'CUcheckpointUnlockArgs' in found_types}}
 
 cdef class CUcheckpointUnlockArgs(CUcheckpointUnlockArgs_st):