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