Merge pull request #116666 from obsidiansystems/cutensor

cudatensor: Init at 1.2.2

maintainers/maintainer-list.nix

@@ -6999,7 +6999,7 @@
     githubId = 3359345;
     name = "obadz";
   };
-  obsidian-systems-maintainence = {
+  obsidian-systems-maintenance = {
     name = "Obsidian Systems Maintenance";
     email = "maintainer@obsidian.systems";
     github = "obsidian-systems-maintenance";

pkgs/development/libraries/science/math/cutensor/default.nix

@@ -0,0 +1,37 @@
+{ callPackage
+, cudatoolkit_10_1, cudatoolkit_10_2
+, cudatoolkit_11_0, cudatoolkit_11_1, cudatoolkit_11_2
+}:
+
+rec {
+  cutensor_cudatoolkit_10_1 = callPackage ./generic.nix rec {
+    version = "1.2.2.5";
+    libPath = "lib/10.1";
+    cudatoolkit = cudatoolkit_10_1;
+    # 1.2.2 is compatible with CUDA 11.0, 11.1, and 11.2:
+    # ephemeral doc at https://developer.nvidia.com/cutensor/downloads
+    sha256 = "1dl9bd71frhac9cb8lvnh71zfsnqxbxbfhndvva2zf6nh0my4klm";
+  };
+
+  cutensor_cudatoolkit_10_2 = cutensor_cudatoolkit_10_1.override {
+    libPath = "lib/10.2";
+    cudatoolkit = cudatoolkit_10_2;
+  };
+
+  cutensor_cudatoolkit_10 = cutensor_cudatoolkit_10_2;
+
+  cutensor_cudatoolkit_11_0 = cutensor_cudatoolkit_10_2.override {
+    libPath = "lib/11";
+    cudatoolkit = cudatoolkit_11_0;
+  };
+
+  cutensor_cudatoolkit_11_1 = cutensor_cudatoolkit_11_0.override {
+    cudatoolkit = cudatoolkit_11_1;
+  };
+
+  cutensor_cudatoolkit_11_2 = cutensor_cudatoolkit_11_0.override {
+    cudatoolkit = cudatoolkit_11_2;
+  };
+
+  cutensor_cudatoolkit_11 = cutensor_cudatoolkit_11_2;
+}

pkgs/development/libraries/science/math/cutensor/generic.nix

@@ -0,0 +1,69 @@
+{ stdenv
+, lib
+, libPath
+, cudatoolkit
+, fetchurl
+, autoPatchelfHook
+, addOpenGLRunpath
+
+, version
+, sha256
+}:
+
+let
+  mostOfVersion = builtins.concatStringsSep "."
+    (lib.take 3 (lib.versions.splitVersion version));
+in
+
+stdenv.mkDerivation {
+  pname = "cudatoolkit-${cudatoolkit.majorVersion}-cutensor";
+  inherit version;
+
+  src = fetchurl {
+    url = "https://developer.download.nvidia.com/compute/cutensor/${mostOfVersion}/local_installers/libcutensor-${stdenv.hostPlatform.parsed.kernel.name}-${stdenv.hostPlatform.parsed.cpu.name}-${version}.tar.gz";
+    inherit sha256;
+  };
+
+  outputs = [ "out" "dev" ];
+
+  nativeBuildInputs = [
+    autoPatchelfHook
+    addOpenGLRunpath
+  ];
+
+  buildInputs = [
+    stdenv.cc.cc.lib
+  ];
+
+  propagatedBuildInputs = [
+    cudatoolkit
+  ];
+
+  # Set RUNPATH so that libcuda in /run/opengl-driver(-32)/lib can be found.
+  # See the explanation in addOpenGLRunpath.
+  installPhase = ''
+    mkdir -p "$out" "$dev"
+    mv include "$dev"
+    mv ${libPath} "$out/lib"
+
+    function finalRPathFixups {
+      for lib in $out/lib/lib*.so; do
+        addOpenGLRunpath $lib
+      done
+    }
+    postFixupHooks+=(finalRPathFixups)
+  '';
+
+  passthru = {
+    inherit cudatoolkit;
+    majorVersion = lib.versions.major version;
+  };
+
+  meta = with lib; {
+    description = "cuTENSOR: A High-Performance CUDA Library For Tensor Primitives";
+    homepage = "https://developer.nvidia.com/cutensor";
+    license = licenses.unfree;
+    platforms = [ "x86_64-linux" ];
+    maintainers = with maintainers; [ obsidian-systems-maintenance ];
+  };
+}

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

@@ -1,7 +1,7 @@
 { lib, buildPythonPackage
 , fetchPypi, isPy3k, linuxPackages
 , fastrlock, numpy, six, wheel, pytest, mock, setuptools
-, cudatoolkit, cudnn, nccl
+, cudatoolkit, cudnn, cutensor, nccl
 }:
 
 buildPythonPackage rec {
@@ -26,6 +26,7 @@ buildPythonPackage rec {
   propagatedBuildInputs = [
     cudatoolkit
     cudnn
+    cutensor
     linuxPackages.nvidia_x11
     nccl
     fastrlock

pkgs/test/cuda/cuda-library-samples/default.nix

@@ -1,16 +1,20 @@
 { callPackage
 , cudatoolkit_10_1, cudatoolkit_10_2
 , cudatoolkit_11_0, cudatoolkit_11_1, cudatoolkit_11_2
+, cutensor_cudatoolkit_10_1, cutensor_cudatoolkit_10_2
+, cutensor_cudatoolkit_11_0, cutensor_cudatoolkit_11_1, cutensor_cudatoolkit_11_2
 }:
 
 rec {
 
   cuda-library-samples_cudatoolkit_10_1 = callPackage ./generic.nix {
     cudatoolkit = cudatoolkit_10_1;
+    cutensor_cudatoolkit = cutensor_cudatoolkit_10_1;
   };
 
   cuda-library-samples_cudatoolkit_10_2 = callPackage ./generic.nix {
     cudatoolkit = cudatoolkit_10_2;
+    cutensor_cudatoolkit = cutensor_cudatoolkit_10_2;
   };
 
   cuda-library-samples_cudatoolkit_10 =
@@ -20,14 +24,17 @@ rec {
 
   cuda-library-samples_cudatoolkit_11_0 = callPackage ./generic.nix {
     cudatoolkit = cudatoolkit_11_0;
+    cutensor_cudatoolkit = cutensor_cudatoolkit_11_0;
   };
 
   cuda-library-samples_cudatoolkit_11_1 = callPackage ./generic.nix {
     cudatoolkit = cudatoolkit_11_1;
+    cutensor_cudatoolkit = cutensor_cudatoolkit_11_1;
   };
 
   cuda-library-samples_cudatoolkit_11_2 = callPackage ./generic.nix {
     cudatoolkit = cudatoolkit_11_2;
+    cutensor_cudatoolkit = cutensor_cudatoolkit_11_2;
   };
 
   cuda-library-samples_cudatoolkit_11 =

pkgs/test/cuda/cuda-library-samples/generic.nix

@@ -1,6 +1,7 @@
 { lib, stdenv, fetchFromGitHub
 , cmake, addOpenGLRunpath
 , cudatoolkit
+, cutensor_cudatoolkit
 }:
 
 let
@@ -29,7 +30,7 @@ let
         cuSPARSE, cuSOLVER, cuFFT, cuRAND, NPP and nvJPEG.
       '';
       license = lib.licenses.bsd3;
-      maintainers = with lib.maintainers; [ obsidian-systems-maintainence ];
+      maintainers = with lib.maintainers; [ obsidian-systems-maintenance ];
     };
   };
 in
@@ -48,4 +49,22 @@ in
 
     sourceRoot = "cuSOLVER/gesv";
   });
+
+  cutensor = stdenv.mkDerivation (commonAttrs // {
+    pname = "cuda-library-samples-cutensor";
+
+    src = "${src}/cuTENSOR";
+
+    cmakeFlags = [
+      "-DCUTENSOR_EXAMPLE_BINARY_INSTALL_DIR=${builtins.placeholder "out"}/bin"
+    ];
+
+    # CUTENSOR_ROOT is double escaped
+    postPatch = ''
+      substituteInPlace CMakeLists.txt \
+        --replace "\''${CUTENSOR_ROOT}/include" "${cutensor_cudatoolkit.dev}/include"
+    '';
+
+    CUTENSOR_ROOT = cutensor_cudatoolkit;
+  });
 }

pkgs/test/cuda/cuda-samples/generic.nix

@@ -46,6 +46,6 @@ stdenv.mkDerivation {
     description = "Samples for CUDA Developers which demonstrates features in CUDA Toolkit";
     # CUDA itself is proprietary, but these sample apps are not.
     license = lib.licenses.bsd3;
-    maintainers = with lib.maintainers; [ obsidian-systems-maintainence ];
+    maintainers = with lib.maintainers; [ obsidian-systems-maintenance ];
   };
 }

pkgs/top-level/all-packages.nix

@@ -3545,6 +3545,18 @@ in
 
   cudnn = cudnn_cudatoolkit_10;
 
+  cutensorPackages = callPackages ../development/libraries/science/math/cutensor { };
+  inherit (cutensorPackages)
+    cutensor_cudatoolkit_10
+    cutensor_cudatoolkit_10_1
+    cutensor_cudatoolkit_10_2
+    cutensor_cudatoolkit_11
+    cutensor_cudatoolkit_11_0
+    cutensor_cudatoolkit_11_1
+    cutensor_cudatoolkit_11_2;
+
+  cutensor = cutensor_cudatoolkit_10;
+
   curlFull = curl.override {
     idnSupport = true;
     ldapSupport = true;

pkgs/top-level/python-packages.nix

@@ -1589,6 +1589,7 @@ in {
     cudatoolkit = pkgs.cudatoolkit_10_0;
     cudnn = pkgs.cudnn_cudatoolkit_10_0;
     nccl = pkgs.nccl_cudatoolkit_10;
+    cutensor = pkgs.cutensor_cudatoolkit_10;
   };
 
   curio = callPackage ../development/python-modules/curio { };