Big nix dependency update (#171)

- Updated tket dependency - This transitively updates the nixpkgs dependency - This necessitated cuda dependency handling changes and version updates - Cupy in particular was awkward to update, but got there in the end with a bit of hackery. --------- Co-authored-by: PabloAndresCQ <[email protected]>
CQCL · Oct 24, 2024 · dd7421f · dd7421f
1 parent 3db8e2b
commit dd7421f
Show file tree

Hide file tree

Showing 8 changed files with 62 additions and 41 deletions.
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -28,6 +28,9 @@
             (self: super: {
               inherit (tket.packages."${system}") tket pytket;
             })
+            (self: super: {
+              mypy' = pkgs.python3Packages.callPackage ./nix-support/mypy.nix {};
+            })
             (self: super: {
               cuda-bundle = pkgs.callPackage ./nix-support/cuda-bundle.nix {};
             })
@@ -49,6 +52,7 @@
         };
         devShells = {
           default = pkgs.mkShell { buildInputs = [ pkgs.pytket-cutensornet ]; };
+          cupy = pkgs.mkShell { buildInputs = [ pkgs.cupy' ]; };
         };
         checks = {
           # no GPU support in checks at the time of writing

diff --git a/nix-support/cuda-bundle.nix b/nix-support/cuda-bundle.nix
@@ -5,27 +5,14 @@
 let
   cutensor' = cudaPackages.callPackage ./cutensor.nix {};
   cuquantum' = cudaPackages.callPackage ./cuquantum.nix {};
-  cusparselt' = cudaPackages.callPackage ./cusparselt.nix {};
 in
 symlinkJoin {
   name = "cuda-bundle-${cudaPackages.cudaVersion}";
   paths = with cudaPackages; [
-    cuda_cccl # <nv/target>
-    cuda_cccl.dev
-    cuda_cudart
-    cuda_nvcc.dev # <crt/host_defines.h>
-    cuda_nvcc
+    cudatoolkit
     cuda_nvprof
-    cuda_nvrtc
-    cuda_nvtx
-    cuda_profiler_api
-    libcublas
-    libcufft
-    libcurand
-    libcusolver
-    libcusparse
-    cusparselt'
-    cusparselt'.dev
+    cuda_cudart.stubs
+    cuda_cudart.static
     cutensor'
     cutensor'.dev
     cuquantum'

diff --git a/nix-support/cupy.nix b/nix-support/cupy.nix
@@ -10,20 +10,20 @@
   mock,
   setuptools,
   cudaPackages,
-  addOpenGLRunpath,
+  addDriverRunpath,
   pythonOlder,
   symlinkJoin,
   substituteAll,
   cuda-bundle
 }:
 
 let
-  inherit (cudaPackages) cudnn nccl;
+  inherit (cudaPackages) nccl;
   cupy-cccl-src = fetchFromGitHub{
     owner = "cupy";
     repo = "cccl";
-    rev = "79ed0e96e35112d171e43f13fa7f324eff7f3de0";
-    hash = "sha256-8YlIp7xOQtB7dKhCL7SIX+cKOvP4V/XPx6EC+Ct4vSc=";
+    rev = "2246a8772a78346d91e580e121cedb820da2648f";
+    hash = "sha256-eLTocR7aKwaPX9NRp+ZpBu/hy2tss3gmJZNDdkSvxS8=";
   };
   dlpack-src = fetchFromGitHub{
     owner = "dmlc";
@@ -40,16 +40,16 @@ let
 in
 buildPythonPackage rec {
   pname = "cupy-cuda12x";
-  version = "13.1.0b";
+  version = "13.3.0";
   format = "setuptools";
 
   disabled = pythonOlder "3.7";
 
   src = fetchFromGitHub {
     owner = "cupy";
     repo = "cupy";
-    rev = "6730353d611f4fd5f6a7494e90c7639a35245337"; # cusparse_lt 0.6.1 support
-    hash = sha256:h6Q6abnq5zPY6mx8dA5l1kXfKjisEOwEqP0rGGu9q0M=;
+    rev = "c2ca5ec9541ff5d754b374ac426846049289bcff";
+    hash = "sha256-BwAnOJuE8IhAv2zEckwq77oIBFMih9j9uY5x16oQCJc=";
   };
   patches = [
     (substituteAll {
@@ -85,13 +85,11 @@ buildPythonPackage rec {
     cuda-bundle
     setuptools
     wheel
-    addOpenGLRunpath
+    addDriverRunpath
     cython_0
   ];
 
   buildInputs = [
-    cuda-bundle
-    cudnn
     nccl
   ];
 
@@ -101,6 +99,7 @@ buildPythonPackage rec {
   LD_LIBRARY_PATH="${cuda-bundle}/lib";
 
   propagatedBuildInputs = [
+    cuda-bundle
     fastrlock
     numpy
   ];
@@ -116,7 +115,7 @@ buildPythonPackage rec {
 
   postFixup = ''
     find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
-      addOpenGLRunpath "$lib"
+      addDriverRunpath "$lib"
       patchelf --set-rpath ${cuda-bundle}/lib $lib;
     done
   '';

diff --git a/nix-support/cupy.patch b/nix-support/cupy.patch
@@ -1,8 +1,20 @@
+diff --git a/cupy/_environment.py b/cupy/_environment.py
+index 105f16b5d..eb36fbc82 100644
+--- a/cupy/_environment.py
++++ b/cupy/_environment.py
+@@ -461,6 +461,7 @@ You can install the library by:
+
+
+ def _get_include_dir_from_conda_or_wheel(major: int, minor: int) -> List[str]:
++    return ["@cuda_bundle@/include"]
+     # FP16 headers from CUDA 12.2+ depends on headers from CUDA Runtime.
+     # See https://github.com/cupy/cupy/issues/8466.
+     if major < 12 or (major == 12 and minor < 2):
 diff --git a/cupy_backends/cuda/libs/_cnvrtc.pxi b/cupy_backends/cuda/libs/_cnvrtc.pxi
-index 9f02b5522..3e0ef9da9 100644
+index 285475036..3e0ef9da9 100644
 --- a/cupy_backends/cuda/libs/_cnvrtc.pxi
 +++ b/cupy_backends/cuda/libs/_cnvrtc.pxi
-@@ -114,30 +114,5 @@ cdef void _initialize() except *:
+@@ -114,33 +114,5 @@ cdef void _initialize() except *:
  cdef SoftLink _get_softlink():
      cdef int runtime_version
      cdef str prefix = 'nvrtc'
@@ -31,6 +43,9 @@ index 9f02b5522..3e0ef9da9 100644
 -        elif runtime_version < 6_00_00000:
 -            # ROCm 5.x
 -            libname = 'libamdhip64.so.5'
+-        elif runtime_version < 7_00_00000:
+-            # ROCm 6.x
+-            libname = 'libamdhip64.so.6'
 -
 +    cdef object libname = '@cuda_bundle@/lib/libnvrtc.so'
      return SoftLink(libname, prefix, mandatory=True)
diff --git a/nix-support/mypy.nix b/nix-support/mypy.nix
@@ -0,0 +1,15 @@
+{
+  python3,
+  fetchFromGitHub,
+  mypy
+}:
+mypy.overrideAttrs (oldAttrs: rec {
+  version = "1.13.0";
+  name = "python${python3.pythonVersion}-mypy-${version}";
+  src = fetchFromGitHub {
+    owner = "python";
+    repo = "mypy";
+    rev = "refs/tags/v${version}";
+    hash = sha256:P2Ozmj7/7QBmjlveHLsNdYgUAerg0qOoa8pO0iQc5os=;
+  };
+})
diff --git a/nix-support/pytket-cutensornet.nix b/nix-support/pytket-cutensornet.nix
@@ -42,7 +42,7 @@ EOF
     propagatedBuildInputs = [ super.pytket super.pycuquantum ];
 
     doCheck = true;
-    checkInputs = with super.python3Packages; [ mypy ];
+    checkInputs = [ super.mypy' ];
     checkPhase = ''
       python -m mypy --config-file=mypy.ini --no-incremental -p pytket
       # running pytest requires GPU access which isn't currently possible (afaik) in the nix sandbox

diff --git a/pytket/extensions/cutensornet/structured_state/mps_mpo.py b/pytket/extensions/cutensornet/structured_state/mps_mpo.py
@@ -550,7 +550,9 @@ def update_variational_tensor(
             optim_fidelity = float(optim_fidelity.real)
 
             # Normalise F and update the variational MPS
-            self._aux_mps.tensors[pos] = F / np.sqrt(optim_fidelity)
+            self._aux_mps.tensors[pos] = F / cp.sqrt(
+                optim_fidelity, dtype=self._cfg._complex_t
+            )
 
             return optim_fidelity