Skip to content

Commit 9ba2616

Browse files
kkraus14kkraus14
committed
Handle CUDA checkpoint restore arg layouts
1 parent cc50515 commit 9ba2616

8 files changed

Lines changed: 315 additions & 198 deletions

File tree

‎cuda_bindings/build_hooks.py‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,13 +109,32 @@ def __init__(self, name, members):
109109
self._name = name
110110
self._member_names = []
111111
self._member_types = []
112+
self._member_declarators = []
112113
for var_name, var_type, _ in members:
113-
var_type = var_type[0]
114-
var_type = var_type.removeprefix("struct ")
115-
var_type = var_type.removeprefix("union ")
114+
base_type = var_type[0]
115+
base_type = base_type.removeprefix("struct ")
116+
base_type = base_type.removeprefix("union ")
116117

117118
self._member_names += [var_name]
118-
self._member_types += [var_type]
119+
self._member_types += [base_type]
120+
self._member_declarators += [tuple(var_type[1:])]
121+
122+
def member_type(self, member_name):
123+
try:
124+
return self._member_types[self._member_names.index(member_name)]
125+
except ValueError:
126+
return None
127+
128+
def member_array_length(self, member_name):
129+
try:
130+
declarators = self._member_declarators[self._member_names.index(member_name)]
131+
except ValueError:
132+
return None
133+
134+
for declarator in declarators:
135+
if isinstance(declarator, list) and len(declarator) == 1:
136+
return declarator[0]
137+
return None
119138

120139
def discoverMembers(self, memberDict, prefix, seen=None):
121140
if seen is None:
@@ -192,6 +211,7 @@ def _parse_headers(header_dict, include_path_list, parser_caching):
192211
# Since we only support 64 bit architectures, we can inline the sizeof(T*) to 8 and then compute the
193212
# result in Python. The arithmetic expression is preserved to help with clarity and understanding
194213
r"char reserved\[52 - sizeof\(CUcheckpointGpuPair \*\)\];": rf"char reserved[{52 - 8}];",
214+
r"char reserved\[64 - sizeof\(CUcheckpointGpuPair \*\) - sizeof\(unsigned int\)\];": rf"char reserved[{64 - 8 - 4}];",
195215
}
196216

197217
print(f'Parsing headers in "{include_path_list}" (Caching = {parser_caching})', flush=True)
@@ -341,6 +361,13 @@ def _build_cuda_bindings(debug=False):
341361
found_types, found_functions, found_values, found_struct, struct_list = _parse_headers(
342362
header_dict, include_path_list, parser_caching
343363
)
364+
struct_field_types = {}
365+
struct_field_array_lengths = {}
366+
for struct_name, struct in struct_list.items():
367+
for member_name in struct._member_names:
368+
key = f"{struct_name}.{member_name}"
369+
struct_field_types[key] = struct.member_type(member_name)
370+
struct_field_array_lengths[key] = struct.member_array_length(member_name)
344371

345372
# Generate code from .in templates
346373
path_list = [
@@ -363,6 +390,8 @@ def _build_cuda_bindings(debug=False):
363390
"found_values": found_values,
364391
"found_struct": found_struct,
365392
"struct_list": struct_list,
393+
"struct_field_types": struct_field_types,
394+
"struct_field_array_lengths": struct_field_array_lengths,
366395
"os": os,
367396
"sys": sys,
368397
"platform": platform,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1719+g565f73f4e. Do not modify it directly.
4+
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1710+g180701c01. Do not modify it directly.
55
from cuda.bindings.cydriver cimport *
66

77
{{if 'cuGetErrorString' in found_functions}}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1719+g565f73f4e. Do not modify it directly.
4+
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1710+g180701c01. Do not modify it directly.
55
{{if 'Windows' == platform.system()}}
66
import os
77
cimport cuda.bindings._lib.windll as windll

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1719+g565f73f4e. Do not modify it directly.
4+
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1710+g180701c01. Do not modify it directly.
55

66
from libc.stdint cimport uint32_t, uint64_t
77

@@ -2458,9 +2458,21 @@ cdef extern from "cuda.h":
24582458
ctypedef CUcheckpointGpuPair_st CUcheckpointGpuPair
24592459

24602460
cdef struct CUcheckpointRestoreArgs_st:
2461+
{{if 'CUcheckpointRestoreArgs_st.gpuPairs' in found_struct}}
24612462
CUcheckpointGpuPair* gpuPairs
2463+
{{endif}}
2464+
{{if 'CUcheckpointRestoreArgs_st.gpuPairsCount' in found_struct}}
24622465
unsigned int gpuPairsCount
2463-
char reserved[52]
2466+
{{endif}}
2467+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
2468+
char reserved[{{struct_field_array_lengths['CUcheckpointRestoreArgs_st.reserved']}}]
2469+
{{endif}}
2470+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
2471+
cuuint64_t reserved[{{struct_field_array_lengths['CUcheckpointRestoreArgs_st.reserved']}}]
2472+
{{endif}}
2473+
{{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
2474+
cuuint64_t reserved1
2475+
{{endif}}
24642476

24652477
ctypedef CUcheckpointRestoreArgs_st CUcheckpointRestoreArgs
24662478

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1719+g565f73f4e. Do not modify it directly.
4+
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1710+g180701c01. Do not modify it directly.
55
cimport cuda.bindings._bindings.cydriver as cydriver
66

77
{{if 'cuGetErrorString' in found_functions}}

‎cuda_bindings/cuda/bindings/driver.pxd.in‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
33

4-
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1719+g565f73f4e. Do not modify it directly.
4+
# This code was automatically generated with version 13.3.0, generator version 0.3.1.dev1710+g180701c01. Do not modify it directly.
55
cimport cuda.bindings.cydriver as cydriver
66

77
include "_lib/utils.pxd"
@@ -5335,10 +5335,18 @@ cdef class CUcheckpointRestoreArgs_st:
53355335
gpuPairsCount : unsigned int
53365336
Number of gpu pairs to remap
53375337
{{endif}}
5338-
{{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
5338+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
53395339
reserved : bytes
53405340
Reserved for future use, must be zeroed
53415341
{{endif}}
5342+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
5343+
reserved : list[cuuint64_t]
5344+
Reserved for future use, must be zeroed
5345+
{{endif}}
5346+
{{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
5347+
reserved1 : cuuint64_t
5348+
Reserved for future use, must be zeroed
5349+
{{endif}}
53425350

53435351
Methods
53445352
-------
@@ -5351,6 +5359,9 @@ cdef class CUcheckpointRestoreArgs_st:
53515359
cdef size_t _gpuPairs_length
53525360
cdef cydriver.CUcheckpointGpuPair* _gpuPairs
53535361
{{endif}}
5362+
{{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
5363+
cdef cuuint64_t _reserved1
5364+
{{endif}}
53545365
{{endif}}
53555366
{{if 'CUcheckpointUnlockArgs_st' in found_struct}}
53565367

@@ -11226,10 +11237,18 @@ cdef class CUcheckpointRestoreArgs(CUcheckpointRestoreArgs_st):
1122611237
gpuPairsCount : unsigned int
1122711238
Number of gpu pairs to remap
1122811239
{{endif}}
11229-
{{if 'CUcheckpointRestoreArgs_st.reserved' in found_struct}}
11240+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'char'}}
1123011241
reserved : bytes
1123111242
Reserved for future use, must be zeroed
1123211243
{{endif}}
11244+
{{if struct_field_types.get('CUcheckpointRestoreArgs_st.reserved') == 'cuuint64_t'}}
11245+
reserved : list[cuuint64_t]
11246+
Reserved for future use, must be zeroed
11247+
{{endif}}
11248+
{{if 'CUcheckpointRestoreArgs_st.reserved1' in found_struct}}
11249+
reserved1 : cuuint64_t
11250+
Reserved for future use, must be zeroed
11251+
{{endif}}
1123311252

1123411253
Methods
1123511254
-------

0 commit comments

Comments
 (0)