python3.pkgs.numba: added optional CUDA support

pkgs/development/python-modules/numba/cuda_path.patch

@@ -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.
+     """

pkgs/development/python-modules/numba/default.nix

@@ -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 = ''

pkgs/top-level/python-packages.nix

@@ -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 { };