Skip to content

Refactor the logic to find CUDA libraries - use try...catch for each library #3575

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
khalatepradnya merged 1 commit into from
Nov 4, 2025
Merged

Refactor the logic to find CUDA libraries - use try...catch for each library #3575

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
220 changes: 126 additions & 94 deletions python/cudaq/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,106 +6,132 @@
# the terms of the Apache License 2.0 which accompanies this distribution. #
# ============================================================================ #

import sys, os, numpy, platform, multiprocessing
from ._packages import get_library_path
import multiprocessing
import os
import sys
import warnings
from pathlib import Path
from typing import Dict, List, Tuple

import numpy

from ._metadata import cuda_major
from ._packages import get_library_path

# Set the multiprocessing start method to 'spawn' if not already set
# Set the multiprocessing start method to `forkserver` if not already set
if multiprocessing.get_start_method(allow_none=True) is None:
multiprocessing.set_start_method('forkserver')


# ============================================================================ #
# CUDA Library Path Configuration
# ============================================================================ #
def _configure_cuda_library_paths() -> None:
"""
Sets the `CUDAQ_DYNLIBS` environment variable with paths to required
CUDA libraries based on the detected CUDA version.
"""
# Skip if already configured or CUDA not detected
if "CUDAQ_DYNLIBS" in os.environ or cuda_major is None:
return

# Library configuration: (package_name_template, library_filename_template)
# Common libraries
common_libs: Dict[str, Tuple[str, str]] = {
'curand': ('nvidia-curand', 'libcurand.so.10'),
'cutensor': ('cutensor-cu{cuda_major}', 'libcutensor.so.2'),
'custatevec': ('custatevec-cu{cuda_major}', 'libcustatevec.so.1'),
'cutensornet': ('cutensornet-cu{cuda_major}', 'libcutensornet.so.2'),
'cudensitymat': ('cudensitymat-cu{cuda_major}', 'libcudensitymat.so.0'),
}

# CUDA 12 specific libraries
cuda_12_specific: Dict[str, Tuple[str, str]] = {
'cudart':
('nvidia-cuda_runtime-cu{cuda_major}', 'libcudart.so.{cuda_major}'),
'nvrtc':
('nvidia-cuda_nvrtc-cu{cuda_major}', 'libnvrtc.so.{cuda_major}'),
'cublas': ('nvidia-cublas-cu{cuda_major}', 'libcublas.so.{cuda_major}'),
'cublaslt':
('nvidia-cublas-cu{cuda_major}', 'libcublasLt.so.{cuda_major}'),
'cusolver': ('nvidia-cusolver-cu{cuda_major}', 'libcusolver.so.11'),
'cusolvermg': ('nvidia-cusolver-cu{cuda_major}', 'libcusolverMg.so.11'),
}

# CUDA 13 specific libraries
cuda_13_specific: Dict[str, Tuple[str, str]] = {
'cudart': ('nvidia-cuda_runtime', 'libcudart.so.{cuda_major}'),
'nvrtc': ('nvidia-cuda_nvrtc', 'libnvrtc.so.{cuda_major}'),
'nvrtc_builtins':
('nvidia-cuda_nvrtc', 'libnvrtc-builtins.so.{cuda_major}.0'),
'cublas': ('nvidia-cublas', 'libcublas.so.{cuda_major}'),
'cublaslt': ('nvidia-cublas', 'libcublasLt.so.{cuda_major}'),
'cusolver': ('nvidia-cusolver', 'libcusolver.so.12'),
'cusolvermg': ('nvidia-cusolver', 'libcusolverMg.so.12'),
}

# Load dependencies first
load_order: List[str] = [
'cudart', 'curand', 'nvrtc', 'nvrtc_builtins', 'cublas', 'cublaslt',
'cusolver', 'cusolvermg', 'cutensor', 'custatevec', 'cutensornet',
'cudensitymat'
]

# Select library configuration based on CUDA version
if cuda_major == 12:
lib_config = {**common_libs, **cuda_12_specific}
elif cuda_major == 13:
lib_config = {**common_libs, **cuda_13_specific}
else:
warnings.warn(f"Unsupported CUDA version {cuda_major}.", RuntimeWarning)
return

# Colon-separated list of library paths for `LinkedLibraryHolder` to load
library_paths: List[str] = []

# Attempt to load each library
for lib_name in load_order:
if lib_name not in lib_config:
continue

package_template, lib_filename_template = lib_config[lib_name]

try:
# Resolve package and library names
package_name = package_template.format(cuda_major=cuda_major)
lib_filename = lib_filename_template.format(cuda_major=cuda_major)
# Get library directory and construct full path
lib_dir = get_library_path(package_name)
lib_path = os.path.join(lib_dir, lib_filename)
# Verify the library file exists
if not os.path.isfile(lib_path):
raise FileNotFoundError(f"Library file not found: {lib_path}")

library_paths.append(lib_path)

except Exception as exc:
warnings.warn(f"Error loading {lib_name}: {exc}.", RuntimeWarning)
continue

os.environ["CUDAQ_DYNLIBS"] = ":".join(library_paths)


# CUDAQ_DYNLIBS must be set before any other imports that would initialize
# LinkedLibraryHolder.
if not "CUDAQ_DYNLIBS" in os.environ and not cuda_major is None:
try:
# common
curand_libs = get_library_path(f"nvidia-curand")
curand_path = os.path.join(curand_libs, "libcurand.so.10")

if cuda_major == 12:
custatevec_libs = get_library_path(f"custatevec-cu{cuda_major}")
custatevec_path = os.path.join(custatevec_libs,
"libcustatevec.so.1")

cutensornet_libs = get_library_path(f"cutensornet-cu{cuda_major}")
cutensornet_path = os.path.join(cutensornet_libs,
"libcutensornet.so.2")

cudensitymat_libs = get_library_path(f"cudensitymat-cu{cuda_major}")
cudensitymat_path = os.path.join(cudensitymat_libs,
"libcudensitymat.so.0")

cutensor_libs = get_library_path(f"cutensor-cu{cuda_major}")
cutensor_path = os.path.join(cutensor_libs, "libcutensor.so.2")

cudart_libs = get_library_path(
f"nvidia-cuda_runtime-cu{cuda_major}")
cudart_path = os.path.join(cudart_libs,
f"libcudart.so.{cuda_major}")

cuda_nvrtc_libs = get_library_path(
f"nvidia-cuda_nvrtc-cu{cuda_major}")
cuda_nvrtc_path = os.path.join(cuda_nvrtc_libs,
f"libnvrtc.so.{cuda_major}")

cusolver_libs = get_library_path(f"nvidia-cusolver-cu{cuda_major}")
cusolver_path = os.path.join(cusolver_libs, f"libcusolver.so.11")
cusolvermg_path = os.path.join(cusolver_libs,
f"libcusolverMg.so.11")

cublas_libs = get_library_path(f"nvidia-cublas-cu{cuda_major}")
cublas_path = os.path.join(cublas_libs,
f"libcublas.so.{cuda_major}")
cublaslt_path = os.path.join(cublas_libs,
f"libcublasLt.so.{cuda_major}")

os.environ[
"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}"
else: # CUDA 13
custatevec_libs = get_library_path(f"custatevec-cu{cuda_major}")
custatevec_path = os.path.join(custatevec_libs,
"libcustatevec.so.1")

cutensornet_libs = get_library_path(f"cutensornet-cu{cuda_major}")
cutensornet_path = os.path.join(cutensornet_libs,
"libcutensornet.so.2")

cudensitymat_libs = get_library_path(f"cudensitymat-cu{cuda_major}")
cudensitymat_path = os.path.join(cudensitymat_libs,
"libcudensitymat.so.0")

cutensor_libs = get_library_path(f"cutensor-cu{cuda_major}")
cutensor_path = os.path.join(cutensor_libs, "libcutensor.so.2")

cudart_libs = get_library_path(f"nvidia-cuda_runtime")
cudart_path = os.path.join(cudart_libs,
f"libcudart.so.{cuda_major}")

cuda_nvrtc_libs = get_library_path(f"nvidia-cuda_nvrtc")
cuda_nvrtc_path = os.path.join(cuda_nvrtc_libs,
f"libnvrtc.so.{cuda_major}")
cuda_nvrtc_builtin_path = os.path.join(
cuda_nvrtc_libs, f"libnvrtc-builtins.so.{cuda_major}.0")

cublas_libs = get_library_path(f"nvidia-cublas")
cublas_path = os.path.join(cublas_libs,
f"libcublas.so.{cuda_major}")
cublaslt_path = os.path.join(cublas_libs,
f"libcublasLt.so.{cuda_major}")

cusolver_libs = get_library_path(f"nvidia-cusolver")
cusolver_path = os.path.join(cusolver_libs, f"libcusolver.so.12")
cusolvermg_path = os.path.join(cusolver_libs,
f"libcusolverMg.so.12")

os.environ[
"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}"
except:
import importlib.util
package_spec = importlib.util.find_spec(f"cuda-quantum-cu{cuda_major}")
if not package_spec is None and not package_spec.loader is None:
print("Could not find a suitable cuQuantum Python package.")
pass
# `LinkedLibraryHolder`.
try:
_configure_cuda_library_paths()
except Exception as exc:
print(f"Error configuring CUDA libraries: {exc}.")

import importlib.util
package_spec = importlib.util.find_spec(f"cuda-quantum-cu{cuda_major}")
if not package_spec is None and not package_spec.loader is None:
print("Could not find a suitable cuQuantum Python package.")
pass

# ============================================================================ #
# Module Imports
# ============================================================================ #

from .display import display_trace
from .kernel.kernel_decorator import kernel, PyKernelDecorator
Expand Down Expand Up @@ -236,6 +262,9 @@
orca = cudaq_runtime.orca


# ============================================================================ #
# Utility Functions
# ============================================================================ #
def synthesize(kernel, *args):
# Compile if necessary, no-op if already compiled
kernel.compile()
Expand Down Expand Up @@ -281,6 +310,9 @@ def __clearKernelRegistries():
from .kernels import uccsd
from .dbg import ast

# ============================================================================ #
# Command Line Argument Parsing
# ============================================================================ #
initKwargs = {}

# Look for --target=<target> options
Expand Down
Loading