about summary refs log tree commit diff
path: root/pkgs/development/libraries/science/math/nccl/default.nix
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/libraries/science/math/nccl/default.nix')
-rw-r--r--pkgs/development/libraries/science/math/nccl/default.nix54
1 files changed, 37 insertions, 17 deletions
diff --git a/pkgs/development/libraries/science/math/nccl/default.nix b/pkgs/development/libraries/science/math/nccl/default.nix
index d3c20258425b..badd08291de4 100644
--- a/pkgs/development/libraries/science/math/nccl/default.nix
+++ b/pkgs/development/libraries/science/math/nccl/default.nix
@@ -1,24 +1,44 @@
-{ callPackage, cudatoolkit_8, cudatoolkit_9 }:
+{ stdenv, fetchFromGitHub, which, cudatoolkit }:
 
-let
-  generic = args: callPackage (import ./generic.nix (removeAttrs args ["cudatoolkit"])) {
-    inherit (args) cudatoolkit;
+stdenv.mkDerivation rec {
+  name = "nccl-${version}-cuda-${cudatoolkit.majorVersion}";
+  version = "2.4.2-1";
+
+  src = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "nccl";
+    rev = "v${version}";
+    sha256 = "0aa4gv51nbmmdhx6vp40l249m4arp30sijrn6kwxdfi1k9kajiq5";
   };
 
-in
+  outputs = [ "out" "dev" ];
 
-{
-  nccl_cudatoolkit_8 = generic rec {
-    version = "2.1.4";
-    cudatoolkit = cudatoolkit_8;
-    srcName = "nccl_${version}-1+cuda${cudatoolkit.majorVersion}_x86_64.txz";
-    sha256 = "1lwwm8kdhna5m318yg304kl2gsz1jwhv4zv4gn8av2m57zh848zi";
-  };
+  nativeBuildInputs = [ which ];
+
+  buildInputs = [ cudatoolkit ];
+
+  preConfigure = ''
+    patchShebangs src/collectives/device/gen_rules.sh
+  '';
+
+  makeFlags = [
+    "CUDA_HOME=${cudatoolkit}"
+    "PREFIX=$(out)"
+  ];
+
+  postFixup = ''
+    moveToOutput lib/libnccl_static.a $dev
+  '';
+
+  NIX_CFLAGS_COMPILE = [ "-Wno-unused-function" ];
+
+  enableParallelBuilding = true;
 
-  nccl_cudatoolkit_9 = generic rec {
-    version = "2.1.4";
-    cudatoolkit = cudatoolkit_9;
-    srcName = "nccl_${version}-1+cuda${cudatoolkit.majorVersion}_x86_64.txz";
-    sha256 = "0pajmqzkacpszs63jh2hw2qqc49kj75kcf7r0ky8hdh560q8xn0p";
+  meta = with stdenv.lib; {
+    description = "Multi-GPU and multi-node collective communication primitives for NVIDIA GPUs";
+    homepage = https://developer.nvidia.com/nccl;
+    license = licenses.bsd3;
+    platforms = [ "x86_64-linux" ];
+    maintainers = with maintainers; [ mdaiter orivej ];
   };
 }