pyrox.dev
1diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py
2index d195bbc29..0699b365b 100644
3--- a/numba/cuda/cuda_paths.py
4+++ b/numba/cuda/cuda_paths.py
5@@ -24,11 +24,7 @@ def _find_valid_path(options):
6
7 def _get_libdevice_path_decision():
8 options = [
9- ('Conda environment', get_conda_ctk()),
10- ('Conda environment (NVIDIA package)', get_nvidia_libdevice_ctk()),
11- ('CUDA_HOME', get_cuda_home('nvvm', 'libdevice')),
12- ('System', get_system_ctk('nvvm', 'libdevice')),
13- ('Debian package', get_debian_pkg_libdevice()),
14+ ('Nix store', get_nix_ctk('nvvm', 'libdevice')),
15 ]
16 by, libdir = _find_valid_path(options)
17 return by, libdir
18@@ -36,17 +32,14 @@ def _get_libdevice_path_decision():
19
20 def _nvvm_lib_dir():
21 if IS_WIN32:
22- return 'nvvm', 'bin'
23+ return 'bin',
24 else:
25- return 'nvvm', 'lib64'
26+ return 'lib64',
27
28
29 def _get_nvvm_path_decision():
30 options = [
31- ('Conda environment', get_conda_ctk()),
32- ('Conda environment (NVIDIA package)', get_nvidia_nvvm_ctk()),
33- ('CUDA_HOME', get_cuda_home(*_nvvm_lib_dir())),
34- ('System', get_system_ctk(*_nvvm_lib_dir())),
35+ ('Nix store', get_nix_ctk(*_nvvm_lib_dir())),
36 ]
37 by, path = _find_valid_path(options)
38 return by, path
39@@ -66,7 +59,7 @@ def _cudalib_path():
40 if IS_WIN32:
41 return 'bin'
42 else:
43- return 'lib64'
44+ return 'lib'
45
46
47 def _cuda_home_static_cudalib_path():
48@@ -78,10 +71,7 @@ def _cuda_home_static_cudalib_path():
49
50 def _get_cudalib_dir_path_decision():
51 options = [
52- ('Conda environment', get_conda_ctk()),
53- ('Conda environment (NVIDIA package)', get_nvidia_cudalib_ctk()),
54- ('CUDA_HOME', get_cuda_home(_cudalib_path())),
55- ('System', get_system_ctk(_cudalib_path())),
56+ ('Nix store', get_nix_lib_ctk(_cudalib_path())),
57 ]
58 by, libdir = _find_valid_path(options)
59 return by, libdir
60@@ -239,3 +229,19 @@ def get_debian_pkg_libdevice():
61 if not os.path.exists(pkg_libdevice_location):
62 return None
63 return pkg_libdevice_location
64+
65+
66+def get_nix_ctk(*subdirs):
67+ """Return path to nix store cudatoolkit; or, None if it doesn't exist.
68+ """
69+ base = '@cuda_toolkit_path@'
70+ if os.path.exists(base):
71+ return os.path.join(base, *subdirs)
72+
73+
74+def get_nix_lib_ctk(*subdirs):
75+ """Return path to nix store cudatoolkit-lib; or, None if it doesn't exist.
76+ """
77+ base = '@cuda_toolkit_lib_path@'
78+ if os.path.exists(base):
79+ return os.path.join(base, *subdirs)