nvidia-thrust: allow omp/tbb instead of cuda

pkgs/development/libraries/nvidia-thrust/default.nix

@@ -1,10 +1,31 @@
 { lib
+, config
 , fetchFromGitHub
 , stdenv
 , cmake
+, pkg-config
+, cudaSupport ? config.cudaSupport or false
 , cudaPackages
 , symlinkJoin
+, enableOpenmp ? true
+, enableTbb ? false
+, tbb
+, hostSystem ? "CPP"
+, deviceSystem ? let
+    devices = lib.optional cudaSupport "CUDA"
+      ++ lib.optional enableTbb "TBB"
+      ++ lib.optional enableOpenmp "OMP"
+      ++ [ "CPP" ];
+  in
+  builtins.head devices
 }:
+
+assert builtins.elem deviceSystem [ "CPP" "OMP" "TBB" "CUDA" ];
+assert builtins.elem hostSystem [ "CPP" "OMP" "TBB" ];
+assert (deviceSystem == "CUDA") -> cudaSupport;
+assert (builtins.elem "TBB" [ deviceSystem hostSystem ]) -> enableTbb;
+assert (builtins.elem "OMP" [ deviceSystem hostSystem ]) -> enableOpenmp;
+
 let
   pname = "nvidia-thrust";
   version = "1.16.0";
@@ -30,12 +51,15 @@ stdenv.mkDerivation {
     hash = "sha256-/EyznxWKuHuvHNjq+SQg27IaRbtkjXR2zlo2YgCWmUQ=";
   };
 
+  buildInputs = lib.optional enableTbb tbb;
+
   nativeBuildInputs = [
     cmake
-
+    pkg-config
+  ] ++ 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
 
@@ -45,9 +69,15 @@ stdenv.mkDerivation {
   ];
 
   cmakeFlags = [
-    "-DTHRUST_INCLUDE_CUB_CMAKE=ON"
+    "-DTHRUST_INCLUDE_CUB_CMAKE=${if cudaSupport then "ON" else "OFF"}"
+    "-DTHRUST_DEVICE_SYSTEM=${deviceSystem}"
+    "-DTHRUST_HOST_SYSTEM=${hostSystem}"
   ];
 
+  passthru = {
+    inherit cudaSupport cudaPackages;
+  };
+
   meta = with lib; {
     description = ''
       A high-level C++ parallel algorithms library

pkgs/top-level/all-packages.nix

@@ -10308,6 +10308,14 @@ with pkgs;
 
   nvidia-thrust = callPackage ../development/libraries/nvidia-thrust { };
 
+  nvidia-thrust-intel = callPackage ../development/libraries/nvidia-thrust {
+    enableTbb = true;
+  };
+
+  nvidia-thrust-cuda = callPackage ../development/libraries/nvidia-thrust {
+    cudaSupport = true;
+  };
+
   miller = callPackage ../tools/text/miller { };
 
   milu = callPackage ../applications/misc/milu { };