22# SPDX-License-Identifier: Apache-2.0
33
44# This module implements basic PEP 517 backend support to defer CUDA-dependent
5- # logic (header parsing, code generation, cythonization) to build time. See:
5+ # logic (cythonization) to build time. See:
66# - https://peps.python.org/pep-0517/
77# - https://setuptools.pypa.io/en/latest/build_meta.html#dynamic-build-dependencies-and-other-build-meta-tweaks
88# - https://github.com/NVIDIA/cuda-python/issues/1635
@@ -203,209 +203,6 @@ def record_build_toolchain() -> None:
203203 record_build_key (_BUILD_TOOLCHAIN_STAMP , name )
204204
205205
206- # -----------------------------------------------------------------------
207- # Header parsing helpers (called only from _build_cuda_bindings)
208-
209- _REQUIRED_HEADERS = {
210- "runtime" : [
211- "driver_types.h" ,
212- "vector_types.h" ,
213- "cuda_runtime.h" ,
214- "surface_types.h" ,
215- "texture_types.h" ,
216- "library_types.h" ,
217- "cuda_runtime_api.h" ,
218- "device_types.h" ,
219- "driver_functions.h" ,
220- "cuda_profiler_api.h" ,
221- ],
222- # nvrtc: headers no longer parsed at build time (pre-generated by cybind).
223- # During compilation, Cython will reference C headers that are not
224- # explicitly parsed above. These are the known dependencies:
225- #
226- # - crt/host_defines.h
227- # - builtin_types.h
228- # - cuda_device_runtime_api.h
229- }
230-
231-
232- class _Struct :
233- def __init__ (self , name , members ):
234- self ._name = name
235- self ._member_names = []
236- self ._member_types = []
237- self ._member_declarators = []
238- for var_name , var_type , _ in members :
239- base_type = var_type [0 ]
240- base_type = base_type .removeprefix ("struct " )
241- base_type = base_type .removeprefix ("union " )
242-
243- self ._member_names += [var_name ]
244- self ._member_types += [base_type ]
245- self ._member_declarators += [tuple (var_type [1 :])]
246-
247- def member_type (self , member_name ):
248- try :
249- return self ._member_types [self ._member_names .index (member_name )]
250- except ValueError :
251- return None
252-
253- def member_array_length (self , member_name ):
254- try :
255- declarators = self ._member_declarators [self ._member_names .index (member_name )]
256- except ValueError :
257- return None
258-
259- for declarator in declarators :
260- if isinstance (declarator , list ) and len (declarator ) == 1 :
261- return declarator [0 ]
262- return None
263-
264- def discoverMembers (self , memberDict , prefix , seen = None ):
265- if seen is None :
266- seen = set ()
267- elif self ._name in seen :
268- return []
269-
270- discovered = []
271- next_seen = set (seen )
272- next_seen .add (self ._name )
273-
274- for memberName , memberType in zip (self ._member_names , self ._member_types ):
275- if memberName :
276- discovered .append ("." .join ([prefix , memberName ]))
277-
278- t = memberType .replace ("const " , "" ).replace ("volatile " , "" ).strip ().rstrip (" *" )
279- if t in memberDict and t != self ._name :
280- discovered += memberDict [t ].discoverMembers (
281- memberDict , discovered [- 1 ] if memberName else prefix , next_seen
282- )
283-
284- return discovered
285-
286- def __repr__ (self ):
287- return f"{ self ._name } : { self ._member_names } with types { self ._member_types } "
288-
289-
290- def _fetch_header_paths (required_headers , include_path_list ):
291- header_dict = {}
292- missing_headers = []
293- for library , header_list in required_headers .items ():
294- header_paths = []
295- for header in header_list :
296- path_candidate = [os .path .join (path , header ) for path in include_path_list ]
297- for path in path_candidate :
298- if os .path .exists (path ):
299- header_paths += [path ]
300- break
301- else :
302- missing_headers += [header ]
303-
304- header_dict [library ] = header_paths
305-
306- if missing_headers :
307- error_message = "Couldn't find required headers: "
308- error_message += ", " .join (missing_headers )
309- cuda_paths = _get_cuda_paths ()
310- raise RuntimeError (f'{ error_message } \n Is CUDA_HOME setup correctly? (CUDA_HOME="{ cuda_paths } ")' )
311-
312- return header_dict
313-
314-
315- def _parse_headers (header_dict , include_path_list , parser_caching ):
316- from pyclibrary import CParser
317-
318- found_types = []
319- found_functions = []
320- found_values = []
321- found_struct = []
322- struct_list = {}
323-
324- replace = {
325- " __device_builtin__ " : " " ,
326- "CUDARTAPI " : " " ,
327- "typedef __device_builtin__ enum cudaError cudaError_t;" : "typedef cudaError cudaError_t;" ,
328- "typedef __device_builtin__ enum cudaOutputMode cudaOutputMode_t;" : "typedef cudaOutputMode cudaOutputMode_t;" ,
329- "typedef enum cudaError cudaError_t;" : "typedef cudaError cudaError_t;" ,
330- "typedef enum cudaOutputMode cudaOutputMode_t;" : "typedef cudaOutputMode cudaOutputMode_t;" ,
331- "typedef enum cudaDataType_t cudaDataType_t;" : "" ,
332- "typedef enum libraryPropertyType_t libraryPropertyType_t;" : "" ,
333- " enum " : " " ,
334- ", enum " : ", " ,
335- "\\ (enum " : "(" ,
336- # Since we only support 64 bit architectures, we can inline the sizeof(T*) to 8 and then compute the
337- # result in Python. The arithmetic expression is preserved to help with clarity and understanding
338- r"char reserved\[52 - sizeof\(CUcheckpointGpuPair \*\)\];" : rf"char reserved[{ 52 - 8 } ];" ,
339- r"char reserved\[64 - sizeof\(CUcheckpointGpuPair \*\) - sizeof\(unsigned int\)\];" : (
340- rf"char reserved[{ 64 - 8 - 4 } ];"
341- ),
342- }
343-
344- print (f'Parsing headers in "{ include_path_list } " (Caching = { parser_caching } )' , flush = True )
345- for library , header_paths in header_dict .items ():
346- print (f"Parsing { library } headers" , flush = True )
347- parser = CParser (
348- header_paths , cache = "./cache_{}" .format (library .split ("." )[0 ]) if parser_caching else None , replace = replace
349- )
350-
351- if library == "driver" :
352- CUDA_VERSION = parser .defs ["macros" ].get ("CUDA_VERSION" , "Unknown" )
353- print (f"Found CUDA_VERSION: { CUDA_VERSION } " , flush = True )
354-
355- found_types += {* parser .defs ["types" ]}
356- found_types += {* parser .defs ["structs" ]}
357- found_types += {* parser .defs ["unions" ]}
358- found_types += {* parser .defs ["enums" ]}
359- found_functions += {* parser .defs ["functions" ]}
360- found_values += {* parser .defs ["values" ]}
361-
362- for key , value in parser .defs ["structs" ].items ():
363- struct_list [key ] = _Struct (key , value ["members" ])
364- for key , value in parser .defs ["unions" ].items ():
365- struct_list [key ] = _Struct (key , value ["members" ])
366-
367- for key , value in struct_list .items ():
368- if key .startswith ("anon_union" ) or key .startswith ("anon_struct" ):
369- continue
370-
371- found_struct += [key ]
372- discovered = value .discoverMembers (struct_list , key )
373- if discovered :
374- found_struct += discovered
375-
376- # TODO(#1312): make this work properly
377- found_types .append ("CUstreamAtomicReductionDataType_enum" )
378-
379- return found_types , found_functions , found_values , found_struct , struct_list
380-
381-
382- # -----------------------------------------------------------------------
383- # Code generation helpers
384-
385-
386- def _fetch_input_files (path ):
387- return [os .path .join (path , f ) for f in os .listdir (path ) if f .endswith (".in" )]
388-
389-
390- def _generate_output (infile , template_vars ):
391- from Cython import Tempita
392-
393- assert infile .endswith (".in" )
394- outfile = infile [:- 3 ]
395-
396- with open (infile , encoding = "utf-8" ) as f :
397- pxdcontent = Tempita .Template (f .read ()).substitute (template_vars )
398-
399- if os .path .exists (outfile ):
400- with open (outfile , encoding = "utf-8" ) as f :
401- if f .read () == pxdcontent :
402- print (f"Skipping { infile } (No change)" , flush = True )
403- return
404- with open (outfile , "w" , encoding = "utf-8" ) as f :
405- print (f"Generating { infile } " , flush = True )
406- f .write (pxdcontent )
407-
408-
409206# -----------------------------------------------------------------------
410207# Extension preparation helpers
411208
@@ -459,9 +256,8 @@ def _prep_extensions(sources, libraries, include_dirs, library_dirs, extra_compi
459256def _build_cuda_bindings (debug = False ):
460257 """Build all cuda-bindings extensions.
461258
462- All CUDA-dependent logic (header parsing, code generation, cythonization)
463- is deferred to this function so that metadata queries do not require a
464- CUDA toolkit installation.
259+ All CUDA-dependent logic (cythonization) is deferred to this function so
260+ that metadata queries do not require a CUDA toolkit installation.
465261 """
466262 import Cython
467263 from Cython .Build import cythonize
@@ -485,7 +281,6 @@ def _build_cuda_bindings(debug=False):
485281 else :
486282 nthreads = int (os .environ .get ("CUDA_PYTHON_PARALLEL_LEVEL" , "0" ) or "0" )
487283
488- parser_caching = bool (os .environ .get ("CUDA_PYTHON_PARSER_CACHING" , False ))
489284 compile_for_coverage = bool (int (os .environ .get ("CUDA_PYTHON_COVERAGE" , "0" )))
490285
491286 toolchain , _cc , _cxx , extra_compile_args , extra_link_args = _resolve_toolchain (
@@ -496,51 +291,8 @@ def _build_cuda_bindings(debug=False):
496291 if debug and sys .platform != "win32" :
497292 extra_cythonize_kwargs ["gdb_debug" ] = True
498293
499- # Parse CUDA headers
500- include_path_list = [os .path .join (path , "include" ) for path in cuda_paths ]
501- header_dict = _fetch_header_paths (_REQUIRED_HEADERS , include_path_list )
502- found_types , found_functions , found_values , found_struct , struct_list = _parse_headers (
503- header_dict , include_path_list , parser_caching
504- )
505- struct_field_types = {}
506- struct_field_array_lengths = {}
507- for struct_name , struct in struct_list .items ():
508- for member_name in struct ._member_names :
509- key = f"{ struct_name } .{ member_name } "
510- struct_field_types [key ] = struct .member_type (member_name )
511- struct_field_array_lengths [key ] = struct .member_array_length (member_name )
512-
513- # Generate code from .in templates
514- path_list = [
515- os .path .join ("cuda" ),
516- os .path .join ("cuda" , "bindings" ),
517- os .path .join ("cuda" , "bindings" , "_bindings" ),
518- os .path .join ("cuda" , "bindings" , "_internal" ),
519- os .path .join ("cuda" , "bindings" , "_lib" ),
520- os .path .join ("cuda" , "bindings" , "utils" ),
521- ]
522- input_files = []
523- for path in path_list :
524- input_files += _fetch_input_files (path )
525-
526- import platform
527-
528- template_vars = {
529- "found_types" : found_types ,
530- "found_functions" : found_functions ,
531- "found_values" : found_values ,
532- "found_struct" : found_struct ,
533- "struct_list" : struct_list ,
534- "struct_field_types" : struct_field_types ,
535- "struct_field_array_lengths" : struct_field_array_lengths ,
536- "os" : os ,
537- "sys" : sys ,
538- "platform" : platform ,
539- }
540- for file in input_files :
541- _generate_output (file , template_vars )
542-
543294 # Prepare compile/link arguments
295+ include_path_list = [os .path .join (path , "include" ) for path in cuda_paths ]
544296 include_dirs = [
545297 os .path .dirname (sysconfig .get_path ("include" )),
546298 ] + include_path_list
@@ -560,14 +312,19 @@ def _cleanup_dst_files():
560312
561313 # Build extension list
562314 extensions = []
563- static_runtime_libraries = ["cudart_static" , "rt" ] if sys .platform == "linux" else ["cudart_static" ]
564315 cuda_bindings_files = glob .glob ("cuda/bindings/*.pyx" )
565316 if sys .platform == "win32" :
566317 cuda_bindings_files = [f for f in cuda_bindings_files if "cufile" not in f ]
318+
319+ def get_static_libraries (f ):
320+ if os .path .basename (f ) in ("runtime.pyx" , "runtime_ptds.pyx" ):
321+ if sys .platform == "linux" :
322+ return ["cudart_static" , "rt" ]
323+ else :
324+ return ["cudart_static" ]
325+ return None
326+
567327 sources_list = [
568- # private
569- (["cuda/bindings/_bindings/cyruntime.pyx" ], static_runtime_libraries ),
570- (["cuda/bindings/_bindings/cyruntime_ptds.pyx" ], static_runtime_libraries ),
571328 # utils
572329 (["cuda/bindings/utils/*.pyx" ], None ),
573330 # public
@@ -576,7 +333,7 @@ def _cleanup_dst_files():
576333 (["cuda/*.pyx" ], None ),
577334 # internal files used by generated bindings
578335 (["cuda/bindings/_internal/utils.pyx" ], None ),
579- * (([f ], None ) for f in dst_files if f .endswith (".pyx" )),
336+ * (([f ], get_static_libraries ( f ) ) for f in dst_files if f .endswith (".pyx" )),
580337 ]
581338
582339 for sources , libraries in sources_list :
0 commit comments