python3Packages.numba: cleanup

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

@@ -1,79 +0,0 @@
-diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py
-index d195bbc29..0699b365b 100644
---- a/numba/cuda/cuda_paths.py
-+++ b/numba/cuda/cuda_paths.py
-@@ -24,11 +24,7 @@ def _find_valid_path(options):
- 
- def _get_libdevice_path_decision():
-     options = [
--        ('Conda environment', get_conda_ctk()),
--        ('Conda environment (NVIDIA package)', get_nvidia_libdevice_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
-@@ -36,17 +32,14 @@ def _get_libdevice_path_decision():
- 
- def _nvvm_lib_dir():
-     if IS_WIN32:
--        return 'nvvm', 'bin'
-+        return 'bin',
-     else:
--        return 'nvvm', 'lib64'
-+        return 'lib64',
- 
- 
- def _get_nvvm_path_decision():
-     options = [
--        ('Conda environment', get_conda_ctk()),
--        ('Conda environment (NVIDIA package)', get_nvidia_nvvm_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
-@@ -66,7 +59,7 @@ def _cudalib_path():
-     if IS_WIN32:
-         return 'bin'
-     else:
--        return 'lib64'
-+        return 'lib'
- 
- 
- def _cuda_home_static_cudalib_path():
-@@ -78,10 +71,7 @@ def _cuda_home_static_cudalib_path():
- 
- def _get_cudalib_dir_path_decision():
-     options = [
--        ('Conda environment', get_conda_ctk()),
--        ('Conda environment (NVIDIA package)', get_nvidia_cudalib_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
-@@ -239,3 +229,19 @@ def get_debian_pkg_libdevice():
-     if not os.path.exists(pkg_libdevice_location):
-         return None
-     return pkg_libdevice_location
-+
-+
-+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)

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

@@ -39,12 +39,27 @@
 }:
 
 let
-  cudatoolkit = cudaPackages.cuda_nvcc;
+  # `libnvvm.so` and `libdevice.10.bc` are bundled in `cuda_nvcc` up to CUDA 12.x, and shipped as a
+  # dedicated redistributable from CUDA 13.0 onwards
+  nvvmRoot =
+    if cudaPackages.cudaOlder "13.0" then
+      "${lib.getLib cudaPackages.cuda_nvcc}/nvvm"
+    else
+      "${lib.getLib cudaPackages.libnvvm}";
+
+  libCudaPath =
+    # Use cuda_compat to provide libcuda.so on pre-Thor Jetsons
+    if (cudaPackages.cuda_compat.meta.available or false) then
+      cudaPackages.cuda_compat
+    # Else, use the host CUDA driver library
+    else
+      addDriverRunpath.driverLink;
 in
 buildPythonPackage (finalAttrs: {
   version = "0.66.0";
   pname = "numba";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "numba";
@@ -65,22 +80,27 @@ buildPythonPackage (finalAttrs: {
     ./numpy2.5.patch
   ]
   ++ lib.optionals cudaSupport [
-    (replaceVars ./cuda_path.patch {
-      cuda_toolkit_path = cudatoolkit;
-      cuda_toolkit_lib_path = lib.getLib cudatoolkit;
+    # Hardcode the paths of the NVIDIA libraries which numba looks up at runtime, instead of
+    # relying on its discovery heuristics (conda environment, `CUDA_HOME`, `/usr/local/cuda`, ...)
+    (replaceVars ./nvidia-libs-paths.patch {
+      libcuda = libCudaPath;
+      libcudart = lib.getLib cudaPackages.cuda_cudart;
+      libcudart_static = lib.getOutput "static" cudaPackages.cuda_cudart;
+      libnvrtc = lib.getLib cudaPackages.cuda_nvrtc;
+      libnvvm = nvvmRoot;
     })
   ];
 
   postPatch = ''
-    substituteInPlace numba/cuda/cudadrv/driver.py \
-      --replace-fail \
-        "dldir = [" \
-        "dldir = [ '${addDriverRunpath.driverLink}/lib', "
-
     substituteInPlace setup.py \
-      --replace-fail 'max_numpy_run_version = "2.5"' 'max_numpy_run_version = "2.6"'
+      --replace-fail \
+        'max_numpy_run_version = "2.5"' \
+        'max_numpy_run_version = "2.6"'
+
     substituteInPlace numba/__init__.py \
-      --replace-fail "(2, 4)" "(2, 6)"
+      --replace-fail \
+        "(2, 4)" \
+        "(2, 6)"
   '';
 
   build-system = [

pkgs/development/python-modules/numba/nvidia-libs-paths.patch

@@ -0,0 +1,43 @@
+diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py
+index 8ff96640d..7eb262a22 100644
+--- a/numba/cuda/cuda_paths.py
++++ b/numba/cuda/cuda_paths.py
+@@ -235,12 +235,20 @@ def get_cuda_paths():
+     if hasattr(get_cuda_paths, '_cached_result'):
+         return get_cuda_paths._cached_result
+     else:
+-        # Not in cache
++        # Not in cache.
++        # The CUDA libraries are installed at fixed, known locations in the Nix
++        # store, so the discovery heuristics above (conda environment,
++        # CUDA_HOME, /usr/local/cuda, ...) are bypassed entirely. Note that
++        # `find_file` also accepts a list of directories, which is needed since
++        # each library lives in its own store path.
+         d = {
+-            'nvvm': _get_nvvm_path(),
+-            'libdevice': _get_libdevice_paths(),
+-            'cudalib_dir': _get_cudalib_dir(),
+-            'static_cudalib_dir': _get_static_cudalib_dir(),
++            'nvvm': _env_path_tuple('Nix store', '@libnvvm@/lib/libnvvm.so'),
++            'libdevice': _env_path_tuple(
++                'Nix store', '@libnvvm@/libdevice/libdevice.10.bc'),
++            'cudalib_dir': _env_path_tuple(
++                'Nix store', ['@libcudart@/lib', '@libnvrtc@/lib']),
++            'static_cudalib_dir': _env_path_tuple(
++                'Nix store', '@libcudart_static@/lib'),
+         }
+         # Cache result
+         get_cuda_paths._cached_result = d
+diff --git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py
+index bea271fce..79312f1ae 100644
+--- a/numba/cuda/cudadrv/driver.py
++++ b/numba/cuda/cudadrv/driver.py
+@@ -118,7 +118,7 @@ def locate_driver_and_loader():
+     else:
+         # Assume to be *nix like
+         dlloader = ctypes.CDLL
+-        dldir = ['/usr/lib', '/usr/lib64']
++        dldir = ['@libcuda@/lib']
+         dlnames = ['libcuda.so', 'libcuda.so.1']
+ 
+     if envpath: