cudatensor: Init at 1.2.2

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-maintainence ];
+  };
+}

pkgs/top-level/all-packages.nix

@@ -3540,6 +3540,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;