ZhengyangZhang's picture
Add files using upload-large-folder tool
28bf3d9 verified
Raw
History Blame Contribute Delete
1.52 kB
from __future__ import annotations
import importlib as _importlib
import os as _os
from cupy_backends.cuda.api import runtime as _runtime
if not _runtime.is_hip:
from cuda import pathfinder as _pathfinder
# All submodules are already made lazily loaded (and they should continue
# remaining so). The problem we are solving is we want to load the required
# C library before the submodule that needs it is imported. This list of
# submodules are supported by the pathfinder.
#
# Note: Do NOT add "nvrtc" here. It is run-time linked instead of dynamically
# linked, so its loading is further deferred (to inside SoftLink).
_submodules = (
'cublas', 'cusolver', 'cusparse', 'curand', 'nccl', 'cutensor',
'cusparselt',
)
_submodules_runtime_link = (
'nvrtc', 'nvtx'
)
def __getattr__(name):
if not _runtime.is_hip and name in _submodules:
try:
pathfinder_name = name
if name == 'cusparselt':
pathfinder_name = 'cusparseLt'
_pathfinder.load_nvidia_dynamic_lib(pathfinder_name)
except _pathfinder.DynamicLibNotFoundError as e:
if (not (_os.environ.get('READTHEDOCS') == 'True') and
not (_os.environ.get('CUPY_CI') is not None)):
raise ImportError(str(e)) from e
if name in (_submodules + _submodules_runtime_link):
return _importlib.import_module(f'cupy_backends.cuda.libs.{name}')
raise AttributeError(
f"module {__name__!r} has no attribute {name!r}")