Backport checkpoint restore arg layout handling to 12.9.x

cuda_bindings/build_hooks.py

@@ -78,13 +78,32 @@ def __init__(self, name, members):
         self._name = name
         self._member_names = []
         self._member_types = []
+        self._member_declarators = []
         for var_name, var_type, _ in members:
-            var_type = var_type[0]
-            var_type = var_type.removeprefix("struct ")
-            var_type = var_type.removeprefix("union ")
+            base_type = var_type[0]
+            base_type = base_type.removeprefix("struct ")
+            base_type = base_type.removeprefix("union ")
 
             self._member_names += [var_name]
-            self._member_types += [var_type]
+            self._member_types += [base_type]
+            self._member_declarators += [tuple(var_type[1:])]
+
+    def member_type(self, member_name):
+        try:
+            return self._member_types[self._member_names.index(member_name)]
+        except ValueError:
+            return None
+
+    def member_array_length(self, member_name):
+        try:
+            declarators = self._member_declarators[self._member_names.index(member_name)]
+        except ValueError:
+            return None
+
+        for declarator in declarators:
+            if isinstance(declarator, list) and len(declarator) == 1:
+                return declarator[0]
+        return None
 
     def discoverMembers(self, memberDict, prefix, seen=None):
         if seen is None:
@@ -161,6 +180,9 @@ def _parse_headers(header_dict, include_path_list, parser_caching):
         # Since we only support 64 bit architectures, we can inline the sizeof(T*) to 8 and then compute the
         # result in Python. The arithmetic expression is preserved to help with clarity and understanding
         r"char reserved\[52 - sizeof\(CUcheckpointGpuPair \*\)\];": rf"char reserved[{52 - 8}];",
+        r"char reserved\[64 - sizeof\(CUcheckpointGpuPair \*\) - sizeof\(unsigned int\)\];": (
+            rf"char reserved[{64 - 8 - 4}];"
+        ),
     }
 
     print(f'Parsing headers in "{include_path_list}" (Caching = {parser_caching})', flush=True)
@@ -310,6 +332,13 @@ def _build_cuda_bindings(strip=False):
     found_types, found_functions, found_values, found_struct, struct_list = _parse_headers(
         header_dict, include_path_list, parser_caching
     )
+    struct_field_types = {}
+    struct_field_array_lengths = {}
+    for struct_name, struct in struct_list.items():
+        for member_name in struct._member_names:
+            key = f"{struct_name}.{member_name}"
+            struct_field_types[key] = struct.member_type(member_name)
+            struct_field_array_lengths[key] = struct.member_array_length(member_name)
 
     # Generate code from .in templates
     path_list = [
@@ -332,6 +361,8 @@ def _build_cuda_bindings(strip=False):
         "found_values": found_values,
         "found_struct": found_struct,
         "struct_list": struct_list,
+        "struct_field_types": struct_field_types,
+        "struct_field_array_lengths": struct_field_array_lengths,
         "os": os,
         "sys": sys,
         "platform": platform,

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

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-# This code was automatically generated with version 12.9.0, generator version 49a8141. Do not modify it directly.
+# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1711+g9f37ce73. Do not modify it directly.
 from cuda.bindings.cydriver cimport *
 
 {{if 'cuGetErrorString' in found_functions}}

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

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1630+gadce055ea.d20260422. Do not modify it directly.
+# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1711+g9f37ce73. Do not modify it directly.
 {{if 'Windows' == platform.system()}}
 import os
 cimport cuda.bindings._lib.windll as windll

cuda_bindings/cuda/bindings/cydriver.pxd.in

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-# This code was automatically generated with version 12.9.0, generator version 49a8141. Do not modify it directly.
+# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1711+g9f37ce73. Do not modify it directly.
 
 from libc.stdint cimport uint32_t, uint64_t
 
@@ -2310,8 +2310,28 @@ cdef extern from "cuda.h":
 
     ctypedef CUcheckpointCheckpointArgs_st CUcheckpointCheckpointArgs
 
+    cdef struct CUcheckpointGpuPair_st:
+        CUuuid oldUuid
+        CUuuid newUuid
+
+    ctypedef CUcheckpointGpuPair_st CUcheckpointGpuPair
+
     cdef struct CUcheckpointRestoreArgs_st:
-        cuuint64_t reserved[8]
+        {{if 'CUcheckpointRestoreArgs_st.gpuPairs' in found_struct}}
+        CUcheckpointGpuPair* gpuPairs
+        {{endif}}
+        {{if 'CUcheckpointRestoreArgs_st.gpuPairsCount' in found_struct}}
+        unsigned int gpuPairsCount
+        {{endif}}
+        {{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
+        char reserved[{{struct_field_array_lengths['CUcheckpointRestoreArgs_st.reserved']}}]
+        {{endif}}
+        {{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
+        cuuint64_t reserved[{{struct_field_array_lengths['CUcheckpointRestoreArgs_st.reserved']}}]
+        {{endif}}
+        {{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
+        cuuint64_t reserved1
+        {{endif}}
 
     ctypedef CUcheckpointRestoreArgs_st CUcheckpointRestoreArgs
 

cuda_bindings/cuda/bindings/cydriver.pyx.in

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-# This code was automatically generated with version 12.9.0, generator version 49a8141. Do not modify it directly.
+# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1711+g9f37ce73. Do not modify it directly.
 cimport cuda.bindings._bindings.cydriver as cydriver
 
 {{if 'cuGetErrorString' in found_functions}}

cuda_bindings/cuda/bindings/driver.pxd.in

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1588+g61faef43a. Do not modify it directly.
+# This code was automatically generated with version 12.9.0, generator version 0.3.1.dev1711+g9f37ce73. Do not modify it directly.
 cimport cuda.bindings.cydriver as cydriver
 
 include "_lib/utils.pxd"
@@ -5089,6 +5089,37 @@ cdef class CUcheckpointCheckpointArgs_st:
     cdef cydriver.CUcheckpointCheckpointArgs_st _pvt_val
     cdef cydriver.CUcheckpointCheckpointArgs_st* _pvt_ptr
 {{endif}}
+{{if 'CUcheckpointGpuPair_st' in found_struct}}
+
+cdef class CUcheckpointGpuPair_st:
+    """
+    CUDA checkpoint GPU UUID pairs for device remapping during restore
+
+    Attributes
+    ----------
+    {{if 'CUcheckpointGpuPair_st.oldUuid' in found_struct}}
+    oldUuid : CUuuid
+        UUID of the GPU that was checkpointed
+    {{endif}}
+    {{if 'CUcheckpointGpuPair_st.newUuid' in found_struct}}
+    newUuid : CUuuid
+        UUID of the GPU to restore onto
+    {{endif}}
+
+    Methods
+    -------
+    getPtr()
+        Get memory address of class instance
+    """
+    cdef cydriver.CUcheckpointGpuPair_st _pvt_val
+    cdef cydriver.CUcheckpointGpuPair_st* _pvt_ptr
+    {{if 'CUcheckpointGpuPair_st.oldUuid' in found_struct}}
+    cdef CUuuid _oldUuid
+    {{endif}}
+    {{if 'CUcheckpointGpuPair_st.newUuid' in found_struct}}
+    cdef CUuuid _newUuid
+    {{endif}}
+{{endif}}
 {{if 'CUcheckpointRestoreArgs_st' in found_struct}}
 
 cdef class CUcheckpointRestoreArgs_st:
@@ -5097,10 +5128,27 @@ cdef class CUcheckpointRestoreArgs_st:
 
     Attributes
     ----------
-    {{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
+    {{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 struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
+    reserved : bytes
+        Reserved for future use, must be zeroed
+    {{endif}}
+    {{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
     reserved : list[cuuint64_t]
         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
     -------
@@ -5109,6 +5157,13 @@ cdef class CUcheckpointRestoreArgs_st:
     """
     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}}
 
@@ -10552,6 +10607,30 @@ cdef class CUcheckpointCheckpointArgs(CUcheckpointCheckpointArgs_st):
     """
     pass
 {{endif}}
+{{if 'CUcheckpointGpuPair' in found_types}}
+
+cdef class CUcheckpointGpuPair(CUcheckpointGpuPair_st):
+    """
+    CUDA checkpoint GPU UUID pairs for device remapping during restore
+
+    Attributes
+    ----------
+    {{if 'CUcheckpointGpuPair_st.oldUuid' in found_struct}}
+    oldUuid : CUuuid
+        UUID of the GPU that was checkpointed
+    {{endif}}
+    {{if 'CUcheckpointGpuPair_st.newUuid' in found_struct}}
+    newUuid : CUuuid
+        UUID of the GPU to restore onto
+    {{endif}}
+
+    Methods
+    -------
+    getPtr()
+        Get memory address of class instance
+    """
+    pass
+{{endif}}
 {{if 'CUcheckpointRestoreArgs' in found_types}}
 
 cdef class CUcheckpointRestoreArgs(CUcheckpointRestoreArgs_st):
@@ -10560,10 +10639,27 @@ cdef class CUcheckpointRestoreArgs(CUcheckpointRestoreArgs_st):
 
     Attributes
     ----------
-    {{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
+    {{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 struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
+    reserved : bytes
+        Reserved for future use, must be zeroed
+    {{endif}}
+    {{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
     reserved : list[cuuint64_t]
         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
     -------