From 05efe64c1bcee8af8d0ae20ca0e3aef164fc8c88 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Wed, 3 Mar 2021 23:53:43 +0000
Subject: [PATCH] cudatensor: Init at 1.2.2

---
 .../science/math/cutensor/default.nix         | 37 ++++++++++
 .../science/math/cutensor/generic.nix         | 69 +++++++++++++++++++
 pkgs/top-level/all-packages.nix               | 12 ++++
 3 files changed, 118 insertions(+)
 create mode 100644 pkgs/development/libraries/science/math/cutensor/default.nix
 create mode 100644 pkgs/development/libraries/science/math/cutensor/generic.nix

diff --git a/pkgs/development/libraries/science/math/cutensor/default.nix b/pkgs/development/libraries/science/math/cutensor/default.nix
new file mode 100644
index 00000000000..1b5895b02bb
--- /dev/null
+++ b/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;
+}
diff --git a/pkgs/development/libraries/science/math/cutensor/generic.nix b/pkgs/development/libraries/science/math/cutensor/generic.nix
new file mode 100644
index 00000000000..31b7dbdb8de
--- /dev/null
+++ b/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 ];
+  };
+}
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index 8b40b3f8877..25e6481698c 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/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;