lol

python3.pkgs.numba: added optional CUDA support

authored by

Ilya Elenskiy and committed by
Frederik Rietdijk
23ae4dfc 5bae60be

+112 -10
+79
pkgs/development/python-modules/numba/cuda_path.patch
··· 1 + diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py 2 + index 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 + """
+25 -9
pkgs/development/python-modules/numba/default.nix
··· 9 9 , llvmlite 10 10 , setuptools 11 11 , libcxx 12 - }: 12 + , substituteAll 13 + 14 + # CUDA-only dependencies: 15 + , addOpenGLRunpath ? null 16 + , cudatoolkit ? null 13 17 18 + # CUDA flags: 19 + , cudaSupport ? false 20 + }: 14 21 buildPythonPackage rec { 15 22 version = "0.55.0"; 16 23 pname = "numba"; ··· 21 28 sha256 = "sha256-siHr2ZdmKh3Ld+TwkUDgIvv+dXetB4H8LgIUE126bL0="; 22 29 }; 23 30 24 - postPatch = '' 25 - substituteInPlace setup.py \ 26 - --replace "1.21" "1.22" 31 + NIX_CFLAGS_COMPILE = lib.optionalString stdenv.isDarwin "-I${lib.getDev libcxx}/include/c++/v1"; 27 32 28 - substituteInPlace numba/__init__.py \ 29 - --replace "(1, 20)" "(1, 21)" 30 - ''; 33 + propagatedBuildInputs = [ numpy llvmlite setuptools ] ++ lib.optionals cudaSupport [ cudatoolkit cudatoolkit.lib ]; 31 34 32 - NIX_CFLAGS_COMPILE = lib.optionalString stdenv.isDarwin "-I${lib.getDev libcxx}/include/c++/v1"; 35 + nativeBuildInputs = lib.optional cudaSupport [ addOpenGLRunpath ]; 36 + 37 + patches = lib.optionals cudaSupport [ 38 + (substituteAll { 39 + src = ./cuda_path.patch; 40 + cuda_toolkit_path = cudatoolkit; 41 + cuda_toolkit_lib_path = cudatoolkit.lib; 42 + }) 43 + ]; 33 44 34 - propagatedBuildInputs = [ numpy llvmlite setuptools ]; 45 + postFixup = lib.optionalString cudaSupport '' 46 + find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do 47 + addOpenGLRunpath "$lib" 48 + patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:$(patchelf --print-rpath "$lib")" "$lib" 49 + done 50 + ''; 35 51 36 52 # Copy test script into $out and run the test suite. 37 53 checkPhase = ''
+8 -1
pkgs/top-level/python-packages.nix
··· 5674 5674 5675 5675 num2words = callPackage ../development/python-modules/num2words { }; 5676 5676 5677 - numba = callPackage ../development/python-modules/numba { }; 5677 + numba = callPackage ../development/python-modules/numba { 5678 + cudaSupport = pkgs.config.cudaSupport or false; 5679 + cudatoolkit = tensorflow_compat_cudatoolkit; 5680 + }; 5681 + 5682 + numbaWithCuda = self.numba.override { 5683 + cudaSupport = true; 5684 + }; 5678 5685 5679 5686 numba-scipy = callPackage ../development/python-modules/numba-scipy { }; 5680 5687