Skip to content

Commit db285ea

Browse files
khalatepradnyakhalatepradnya
authored
* Refactor the logic to find CUDA libraries - use try...catch for each library. (#3575)
Signed-off-by: Pradnya Khalate <[email protected]>
1 parent a09b122 commit db285ea

File tree

1 file changed

+126
-94
lines changed

1 file changed

+126
-94
lines changed

‎python/cudaq/__init__.py‎

Lines changed: 126 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -6,106 +6,132 @@
66
# the terms of the Apache License 2.0 which accompanies this distribution. #
77
# ============================================================================ #
88

9-
import sys, os, numpy, platform, multiprocessing
10-
from ._packages import get_library_path
9+
import multiprocessing
10+
import os
11+
import sys
12+
import warnings
13+
from pathlib import Path
14+
from typing import Dict, List, Tuple
15+
16+
import numpy
17+
1118
from ._metadata import cuda_major
19+
from ._packages import get_library_path
1220

13-
# Set the multiprocessing start method to 'spawn' if not already set
21+
# Set the multiprocessing start method to `forkserver` if not already set
1422
if multiprocessing.get_start_method(allow_none=True) is None:
1523
multiprocessing.set_start_method('forkserver')
1624

25+
26+
# ============================================================================ #
27+
# CUDA Library Path Configuration
28+
# ============================================================================ #
29+
def _configure_cuda_library_paths() -> None:
30+
"""
31+
Sets the `CUDAQ_DYNLIBS` environment variable with paths to required
32+
CUDA libraries based on the detected CUDA version.
33+
"""
34+
# Skip if already configured or CUDA not detected
35+
if "CUDAQ_DYNLIBS" in os.environ or cuda_major is None:
36+
return
37+
38+
# Library configuration: (package_name_template, library_filename_template)
39+
# Common libraries
40+
common_libs: Dict[str, Tuple[str, str]] = {
41+
'curand': ('nvidia-curand', 'libcurand.so.10'),
42+
'cutensor': ('cutensor-cu{cuda_major}', 'libcutensor.so.2'),
43+
'custatevec': ('custatevec-cu{cuda_major}', 'libcustatevec.so.1'),
44+
'cutensornet': ('cutensornet-cu{cuda_major}', 'libcutensornet.so.2'),
45+
'cudensitymat': ('cudensitymat-cu{cuda_major}', 'libcudensitymat.so.0'),
46+
}
47+
48+
# CUDA 12 specific libraries
49+
cuda_12_specific: Dict[str, Tuple[str, str]] = {
50+
'cudart':
51+
('nvidia-cuda_runtime-cu{cuda_major}', 'libcudart.so.{cuda_major}'),
52+
'nvrtc':
53+
('nvidia-cuda_nvrtc-cu{cuda_major}', 'libnvrtc.so.{cuda_major}'),
54+
'cublas': ('nvidia-cublas-cu{cuda_major}', 'libcublas.so.{cuda_major}'),
55+
'cublaslt':
56+
('nvidia-cublas-cu{cuda_major}', 'libcublasLt.so.{cuda_major}'),
57+
'cusolver': ('nvidia-cusolver-cu{cuda_major}', 'libcusolver.so.11'),
58+
'cusolvermg': ('nvidia-cusolver-cu{cuda_major}', 'libcusolverMg.so.11'),
59+
}
60+
61+
# CUDA 13 specific libraries
62+
cuda_13_specific: Dict[str, Tuple[str, str]] = {
63+
'cudart': ('nvidia-cuda_runtime', 'libcudart.so.{cuda_major}'),
64+
'nvrtc': ('nvidia-cuda_nvrtc', 'libnvrtc.so.{cuda_major}'),
65+
'nvrtc_builtins':
66+
('nvidia-cuda_nvrtc', 'libnvrtc-builtins.so.{cuda_major}.0'),
67+
'cublas': ('nvidia-cublas', 'libcublas.so.{cuda_major}'),
68+
'cublaslt': ('nvidia-cublas', 'libcublasLt.so.{cuda_major}'),
69+
'cusolver': ('nvidia-cusolver', 'libcusolver.so.12'),
70+
'cusolvermg': ('nvidia-cusolver', 'libcusolverMg.so.12'),
71+
}
72+
73+
# Load dependencies first
74+
load_order: List[str] = [
75+
'cudart', 'curand', 'nvrtc', 'nvrtc_builtins', 'cublas', 'cublaslt',
76+
'cusolver', 'cusolvermg', 'cutensor', 'custatevec', 'cutensornet',
77+
'cudensitymat'
78+
]
79+
80+
# Select library configuration based on CUDA version
81+
if cuda_major == 12:
82+
lib_config = {**common_libs, **cuda_12_specific}
83+
elif cuda_major == 13:
84+
lib_config = {**common_libs, **cuda_13_specific}
85+
else:
86+
warnings.warn(f"Unsupported CUDA version {cuda_major}.", RuntimeWarning)
87+
return
88+
89+
# Colon-separated list of library paths for `LinkedLibraryHolder` to load
90+
library_paths: List[str] = []
91+
92+
# Attempt to load each library
93+
for lib_name in load_order:
94+
if lib_name not in lib_config:
95+
continue
96+
97+
package_template, lib_filename_template = lib_config[lib_name]
98+
99+
try:
100+
# Resolve package and library names
101+
package_name = package_template.format(cuda_major=cuda_major)
102+
lib_filename = lib_filename_template.format(cuda_major=cuda_major)
103+
# Get library directory and construct full path
104+
lib_dir = get_library_path(package_name)
105+
lib_path = os.path.join(lib_dir, lib_filename)
106+
# Verify the library file exists
107+
if not os.path.isfile(lib_path):
108+
raise FileNotFoundError(f"Library file not found: {lib_path}")
109+
110+
library_paths.append(lib_path)
111+
112+
except Exception as exc:
113+
warnings.warn(f"Error loading {lib_name}: {exc}.", RuntimeWarning)
114+
continue
115+
116+
os.environ["CUDAQ_DYNLIBS"] = ":".join(library_paths)
117+
118+
17119
# CUDAQ_DYNLIBS must be set before any other imports that would initialize
18-
# LinkedLibraryHolder.
19-
if not "CUDAQ_DYNLIBS" in os.environ and not cuda_major is None:
20-
try:
21-
# common
22-
curand_libs = get_library_path(f"nvidia-curand")
23-
curand_path = os.path.join(curand_libs, "libcurand.so.10")
24-
25-
if cuda_major == 12:
26-
custatevec_libs = get_library_path(f"custatevec-cu{cuda_major}")
27-
custatevec_path = os.path.join(custatevec_libs,
28-
"libcustatevec.so.1")
29-
30-
cutensornet_libs = get_library_path(f"cutensornet-cu{cuda_major}")
31-
cutensornet_path = os.path.join(cutensornet_libs,
32-
"libcutensornet.so.2")
33-
34-
cudensitymat_libs = get_library_path(f"cudensitymat-cu{cuda_major}")
35-
cudensitymat_path = os.path.join(cudensitymat_libs,
36-
"libcudensitymat.so.0")
37-
38-
cutensor_libs = get_library_path(f"cutensor-cu{cuda_major}")
39-
cutensor_path = os.path.join(cutensor_libs, "libcutensor.so.2")
40-
41-
cudart_libs = get_library_path(
42-
f"nvidia-cuda_runtime-cu{cuda_major}")
43-
cudart_path = os.path.join(cudart_libs,
44-
f"libcudart.so.{cuda_major}")
45-
46-
cuda_nvrtc_libs = get_library_path(
47-
f"nvidia-cuda_nvrtc-cu{cuda_major}")
48-
cuda_nvrtc_path = os.path.join(cuda_nvrtc_libs,
49-
f"libnvrtc.so.{cuda_major}")
50-
51-
cusolver_libs = get_library_path(f"nvidia-cusolver-cu{cuda_major}")
52-
cusolver_path = os.path.join(cusolver_libs, f"libcusolver.so.11")
53-
cusolvermg_path = os.path.join(cusolver_libs,
54-
f"libcusolverMg.so.11")
55-
56-
cublas_libs = get_library_path(f"nvidia-cublas-cu{cuda_major}")
57-
cublas_path = os.path.join(cublas_libs,
58-
f"libcublas.so.{cuda_major}")
59-
cublaslt_path = os.path.join(cublas_libs,
60-
f"libcublasLt.so.{cuda_major}")
61-
62-
os.environ[
63-
"CUDAQ_DYNLIBS"] = f"{cudart_path}:{curand_path}:{cuda_nvrtc_path}:{cublaslt_path}:{cublas_path}:{cusolver_path}:{cusolvermg_path}:{cutensor_path}:{custatevec_path}:{cutensornet_path}:{cudensitymat_path}"
64-
else: # CUDA 13
65-
custatevec_libs = get_library_path(f"custatevec-cu{cuda_major}")
66-
custatevec_path = os.path.join(custatevec_libs,
67-
"libcustatevec.so.1")
68-
69-
cutensornet_libs = get_library_path(f"cutensornet-cu{cuda_major}")
70-
cutensornet_path = os.path.join(cutensornet_libs,
71-
"libcutensornet.so.2")
72-
73-
cudensitymat_libs = get_library_path(f"cudensitymat-cu{cuda_major}")
74-
cudensitymat_path = os.path.join(cudensitymat_libs,
75-
"libcudensitymat.so.0")
76-
77-
cutensor_libs = get_library_path(f"cutensor-cu{cuda_major}")
78-
cutensor_path = os.path.join(cutensor_libs, "libcutensor.so.2")
79-
80-
cudart_libs = get_library_path(f"nvidia-cuda_runtime")
81-
cudart_path = os.path.join(cudart_libs,
82-
f"libcudart.so.{cuda_major}")
83-
84-
cuda_nvrtc_libs = get_library_path(f"nvidia-cuda_nvrtc")
85-
cuda_nvrtc_path = os.path.join(cuda_nvrtc_libs,
86-
f"libnvrtc.so.{cuda_major}")
87-
cuda_nvrtc_builtin_path = os.path.join(
88-
cuda_nvrtc_libs, f"libnvrtc-builtins.so.{cuda_major}.0")
89-
90-
cublas_libs = get_library_path(f"nvidia-cublas")
91-
cublas_path = os.path.join(cublas_libs,
92-
f"libcublas.so.{cuda_major}")
93-
cublaslt_path = os.path.join(cublas_libs,
94-
f"libcublasLt.so.{cuda_major}")
95-
96-
cusolver_libs = get_library_path(f"nvidia-cusolver")
97-
cusolver_path = os.path.join(cusolver_libs, f"libcusolver.so.12")
98-
cusolvermg_path = os.path.join(cusolver_libs,
99-
f"libcusolverMg.so.12")
100-
101-
os.environ[
102-
"CUDAQ_DYNLIBS"] = f"{cudart_path}:{curand_path}:{cuda_nvrtc_path}:{cuda_nvrtc_builtin_path}:{cublas_path}:{cublaslt_path}:{cusolver_path}:{cusolvermg_path}:{cutensor_path}:{custatevec_path}:{cutensornet_path}:{cudensitymat_path}"
103-
except:
104-
import importlib.util
105-
package_spec = importlib.util.find_spec(f"cuda-quantum-cu{cuda_major}")
106-
if not package_spec is None and not package_spec.loader is None:
107-
print("Could not find a suitable cuQuantum Python package.")
108-
pass
120+
# `LinkedLibraryHolder`.
121+
try:
122+
_configure_cuda_library_paths()
123+
except Exception as exc:
124+
print(f"Error configuring CUDA libraries: {exc}.")
125+
126+
import importlib.util
127+
package_spec = importlib.util.find_spec(f"cuda-quantum-cu{cuda_major}")
128+
if not package_spec is None and not package_spec.loader is None:
129+
print("Could not find a suitable cuQuantum Python package.")
130+
pass
131+
132+
# ============================================================================ #
133+
# Module Imports
134+
# ============================================================================ #
109135

110136
from .display import display_trace
111137
from .kernel.kernel_decorator import kernel, PyKernelDecorator
@@ -236,6 +262,9 @@
236262
orca = cudaq_runtime.orca
237263

238264

265+
# ============================================================================ #
266+
# Utility Functions
267+
# ============================================================================ #
239268
def synthesize(kernel, *args):
240269
# Compile if necessary, no-op if already compiled
241270
kernel.compile()
@@ -281,6 +310,9 @@ def __clearKernelRegistries():
281310
from .kernels import uccsd
282311
from .dbg import ast
283312

313+
# ============================================================================ #
314+
# Command Line Argument Parsing
315+
# ============================================================================ #
284316
initKwargs = {}
285317

286318
# Look for --target=<target> options

0 commit comments

Comments
 (0)