nvidia-thrust: cudatoolkit -> redist cudaPackages

1 files changed, 24 insertions, 10 deletions

diff --git a/pkgs/development/libraries/nvidia-thrust/default.nix b/pkgs/development/libraries/nvidia-thrust/default.nix
index 545607ead2d..679b2c61c45 100644
--- a/pkgs/development/libraries/nvidia-thrust/default.nix
+++ b/pkgs/development/libraries/nvidia-thrust/default.nix
@@ -4,7 +4,7 @@
 , stdenv
 , cmake
 , pkg-config
-, cudaPackages
+, cudaPackages ? { }
 , symlinkJoin
 , tbb
 , hostSystem ? "CPP"
@@ -27,6 +27,9 @@ let
   pname = "nvidia-thrust";
   version = "1.16.0";
 
+  inherit (cudaPackages) backendStdenv cudaFlags;
+  cudaCapabilities = map cudaFlags.dropDot cudaFlags.cudaCapabilities;
+
   tbbSupport = builtins.elem "TBB" [ deviceSystem hostSystem ];
   cudaSupport = deviceSystem == "CUDA";
 
@@ -35,9 +38,13 @@ let
     name = "cuda-packages-unsplit";
     paths = with cudaPackages; [
       cuda_nvcc
+      cuda_nvrtc # symbols: cudaLaunchDevice, &c; notice postBuild
       cuda_cudart # cuda_runtime.h
       libcublas
     ];
+    postBuild = ''
+      ln -s $out/lib $out/lib64
+    '';
   };
 in
 stdenv.mkDerivation {
@@ -51,6 +58,15 @@ stdenv.mkDerivation {
     hash = "sha256-/EyznxWKuHuvHNjq+SQg27IaRbtkjXR2zlo2YgCWmUQ=";
   };
 
+  # NVIDIA's "compiler hacks" seem like work-arounds for legacy toolchains and
+  # cause us errors such as:
+  # > Thrust's test harness uses CMAKE_CXX_COMPILER for the CUDA host compiler.
+  # > Refusing to overwrite specified CMAKE_CUDA_HOST_COMPILER
+  # So we un-fix cmake after them:
+  postPatch = ''
+    echo > cmake/ThrustCompilerHacks.cmake
+  '';
+
   buildInputs = lib.optionals tbbSupport [ tbb ];
 
   nativeBuildInputs = [
@@ -59,23 +75,21 @@ stdenv.mkDerivation {
   ] ++ lib.optionals cudaSupport [
     # Goes in native build inputs because thrust looks for headers
     # in a path relative to nvcc...
-
-    # Works when build=host, but we only support
-    # cuda on x86_64 anyway
-
-    # TODO: but instead using this
-    # cudaJoined
-    cudaPackages.cudatoolkit
+    cudaJoined
   ];
 
   cmakeFlags = [
     "-DTHRUST_INCLUDE_CUB_CMAKE=${if cudaSupport then "ON" else "OFF"}"
     "-DTHRUST_DEVICE_SYSTEM=${deviceSystem}"
     "-DTHRUST_HOST_SYSTEM=${hostSystem}"
-  ];
+    "-DTHRUST_AUTO_DETECT_COMPUTE_ARCHS=OFF"
+    "-DTHRUST_DISABLE_ARCH_BY_DEFAULT=ON"
+  ] ++ lib.optionals cudaFlags.enableForwardCompat [
+    "-DTHRUST_ENABLE_COMPUTE_FUTURE=ON"
+  ] ++ map (sm: "THRUST_ENABLE_COMPUTE_${sm}") cudaCapabilities;
 
   passthru = {
-    inherit cudaSupport cudaPackages;
+    inherit cudaSupport cudaPackages cudaJoined;
   };
 
   meta = with lib; {