python3.pkgs.numba: added optional CUDA support
This commit is contained in:
parent
5bae60beee
commit
23ae4dfc22
79
pkgs/development/python-modules/numba/cuda_path.patch
Normal file
79
pkgs/development/python-modules/numba/cuda_path.patch
Normal file
|
@ -0,0 +1,79 @@
|
|||
diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py
|
||||
index b9988bc..a642680 100644
|
||||
--- a/numba/cuda/cuda_paths.py
|
||||
+++ b/numba/cuda/cuda_paths.py
|
||||
@@ -24,10 +24,7 @@ def _find_valid_path(options):
|
||||
|
||||
def _get_libdevice_path_decision():
|
||||
options = [
|
||||
- ('Conda environment', get_conda_ctk()),
|
||||
- ('CUDA_HOME', get_cuda_home('nvvm', 'libdevice')),
|
||||
- ('System', get_system_ctk('nvvm', 'libdevice')),
|
||||
- ('Debian package', get_debian_pkg_libdevice()),
|
||||
+ ('Nix store', get_nix_ctk('nvvm', 'libdevice')),
|
||||
]
|
||||
by, libdir = _find_valid_path(options)
|
||||
return by, libdir
|
||||
@@ -35,18 +32,16 @@ def _get_libdevice_path_decision():
|
||||
|
||||
def _nvvm_lib_dir():
|
||||
if IS_WIN32:
|
||||
- return 'nvvm', 'bin'
|
||||
+ return 'bin',
|
||||
elif IS_OSX:
|
||||
- return 'nvvm', 'lib'
|
||||
+ return 'lib',
|
||||
else:
|
||||
- return 'nvvm', 'lib64'
|
||||
+ return 'lib64',
|
||||
|
||||
|
||||
def _get_nvvm_path_decision():
|
||||
options = [
|
||||
- ('Conda environment', get_conda_ctk()),
|
||||
- ('CUDA_HOME', get_cuda_home(*_nvvm_lib_dir())),
|
||||
- ('System', get_system_ctk(*_nvvm_lib_dir())),
|
||||
+ ('Nix store', get_nix_ctk(*_nvvm_lib_dir())),
|
||||
]
|
||||
by, path = _find_valid_path(options)
|
||||
return by, path
|
||||
@@ -74,14 +69,12 @@ def _cudalib_path():
|
||||
elif IS_OSX:
|
||||
return 'lib'
|
||||
else:
|
||||
- return 'lib64'
|
||||
+ return 'lib'
|
||||
|
||||
|
||||
def _get_cudalib_dir_path_decision():
|
||||
options = [
|
||||
- ('Conda environment', get_conda_ctk()),
|
||||
- ('CUDA_HOME', get_cuda_home(_cudalib_path())),
|
||||
- ('System', get_system_ctk(_cudalib_path())),
|
||||
+ ('Nix store', get_nix_lib_ctk(_cudalib_path())),
|
||||
]
|
||||
by, libdir = _find_valid_path(options)
|
||||
return by, libdir
|
||||
@@ -92,6 +85,22 @@ def _get_cudalib_dir():
|
||||
return _env_path_tuple(by, libdir)
|
||||
|
||||
|
||||
+def get_nix_ctk(*subdirs):
|
||||
+ """Return path to nix store cudatoolkit; or, None if it doesn't exist.
|
||||
+ """
|
||||
+ base = '@cuda_toolkit_path@'
|
||||
+ if os.path.exists(base):
|
||||
+ return os.path.join(base, *subdirs)
|
||||
+
|
||||
+
|
||||
+def get_nix_lib_ctk(*subdirs):
|
||||
+ """Return path to nix store cudatoolkit-lib; or, None if it doesn't exist.
|
||||
+ """
|
||||
+ base = '@cuda_toolkit_lib_path@'
|
||||
+ if os.path.exists(base):
|
||||
+ return os.path.join(base, *subdirs)
|
||||
+
|
||||
+
|
||||
def get_system_ctk(*subdirs):
|
||||
"""Return path to system-wide cudatoolkit; or, None if it doesn't exist.
|
||||
"""
|
|
@ -9,8 +9,15 @@
|
|||
, llvmlite
|
||||
, setuptools
|
||||
, libcxx
|
||||
}:
|
||||
, substituteAll
|
||||
|
||||
# CUDA-only dependencies:
|
||||
, addOpenGLRunpath ? null
|
||||
, cudatoolkit ? null
|
||||
|
||||
# CUDA flags:
|
||||
, cudaSupport ? false
|
||||
}:
|
||||
buildPythonPackage rec {
|
||||
version = "0.55.0";
|
||||
pname = "numba";
|
||||
|
@ -21,17 +28,26 @@ buildPythonPackage rec {
|
|||
sha256 = "sha256-siHr2ZdmKh3Ld+TwkUDgIvv+dXetB4H8LgIUE126bL0=";
|
||||
};
|
||||
|
||||
postPatch = ''
|
||||
substituteInPlace setup.py \
|
||||
--replace "1.21" "1.22"
|
||||
|
||||
substituteInPlace numba/__init__.py \
|
||||
--replace "(1, 20)" "(1, 21)"
|
||||
'';
|
||||
|
||||
NIX_CFLAGS_COMPILE = lib.optionalString stdenv.isDarwin "-I${lib.getDev libcxx}/include/c++/v1";
|
||||
|
||||
propagatedBuildInputs = [ numpy llvmlite setuptools ];
|
||||
propagatedBuildInputs = [ numpy llvmlite setuptools ] ++ lib.optionals cudaSupport [ cudatoolkit cudatoolkit.lib ];
|
||||
|
||||
nativeBuildInputs = lib.optional cudaSupport [ addOpenGLRunpath ];
|
||||
|
||||
patches = lib.optionals cudaSupport [
|
||||
(substituteAll {
|
||||
src = ./cuda_path.patch;
|
||||
cuda_toolkit_path = cudatoolkit;
|
||||
cuda_toolkit_lib_path = cudatoolkit.lib;
|
||||
})
|
||||
];
|
||||
|
||||
postFixup = lib.optionalString cudaSupport ''
|
||||
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
|
||||
addOpenGLRunpath "$lib"
|
||||
patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:$(patchelf --print-rpath "$lib")" "$lib"
|
||||
done
|
||||
'';
|
||||
|
||||
# Copy test script into $out and run the test suite.
|
||||
checkPhase = ''
|
||||
|
|
|
@ -5674,7 +5674,14 @@ in {
|
|||
|
||||
num2words = callPackage ../development/python-modules/num2words { };
|
||||
|
||||
numba = callPackage ../development/python-modules/numba { };
|
||||
numba = callPackage ../development/python-modules/numba {
|
||||
cudaSupport = pkgs.config.cudaSupport or false;
|
||||
cudatoolkit = tensorflow_compat_cudatoolkit;
|
||||
};
|
||||
|
||||
numbaWithCuda = self.numba.override {
|
||||
cudaSupport = true;
|
||||
};
|
||||
|
||||
numba-scipy = callPackage ../development/python-modules/numba-scipy { };
|
||||
|
||||
|
|
Loading…
Reference in a new issue