From dd7421fb96f449a760f3bab0afc1da9223f2c4e8 Mon Sep 17 00:00:00 2001 From: Jake Arkinstall <65358059+jake-arkinstall@users.noreply.github.com> Date: Thu, 24 Oct 2024 17:51:24 +0100 Subject: [PATCH] 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 --- flake.lock | 19 ++++++++--------- flake.nix | 4 ++++ nix-support/cuda-bundle.nix | 19 +++-------------- nix-support/cupy.nix | 21 +++++++++---------- nix-support/cupy.patch | 19 +++++++++++++++-- nix-support/mypy.nix | 15 +++++++++++++ nix-support/pytket-cutensornet.nix | 2 +- .../cutensornet/structured_state/mps_mpo.py | 4 +++- 8 files changed, 62 insertions(+), 41 deletions(-) create mode 100644 nix-support/mypy.nix diff --git a/flake.lock b/flake.lock index 78253789..a2df5b08 100644 --- a/flake.lock +++ b/flake.lock @@ -38,11 +38,11 @@ "systems": "systems_2" }, "locked": { - "lastModified": 1710146030, - "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=", + "lastModified": 1726560853, + "narHash": "sha256-X6rJYSESBVr3hBoH0WbKE5KvhPU5bloyZ2L4K60/fPQ=", "owner": "numtide", "repo": "flake-utils", - "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a", + "rev": "c1dfcf08411b08f6b8615f7d8971a2bfa81d5e8a", "type": "github" }, "original": { @@ -87,16 +87,15 @@ }, "nixpkgs_2": { "locked": { - "lastModified": 1722087241, - "narHash": "sha256-2ShmEaFi0kJVOEEu5gmlykN5dwjWYWYUJmlRTvZQRpU=", + "lastModified": 1729693369, + "narHash": "sha256-45rettjkIaKhKLEK3KDYaxWIJstCd6G+W8idXXu3+YM=", "owner": "nixos", "repo": "nixpkgs", - "rev": "8c50662509100d53229d4be607f1a3a31157fa12", + "rev": "9a79bc998d0207268c8e00b737e896e2e717ccd2", "type": "github" }, "original": { "owner": "nixos", - "ref": "nixos-24.05", "repo": "nixpkgs", "type": "github" } @@ -148,11 +147,11 @@ "nixpkgs": "nixpkgs_2" }, "locked": { - "lastModified": 1722327570, - "narHash": "sha256-e4T0mMg+zGd+PlaY3OXOGQ6gclbwjVoapRlnHJ6gDf0=", + "lastModified": 1729759304, + "narHash": "sha256-nBdv1glaLmuHLdb66CitZ+3NM+UTeZBxipQFQopoOMc=", "owner": "CQCL", "repo": "tket", - "rev": "da1f491059b608701701da6202abb96c67ac6031", + "rev": "7bae471e832aca0fba5910e5080640a4b9a2dfe2", "type": "github" }, "original": { diff --git a/flake.nix b/flake.nix index 72eb5a54..f358dc8b 100644 --- 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 index ad2ad430..68e26f58 100644 --- 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 # - cuda_cccl.dev - cuda_cudart - cuda_nvcc.dev # - 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 index 932a7cb5..cf938a89 100644 --- a/nix-support/cupy.nix +++ b/nix-support/cupy.nix @@ -10,7 +10,7 @@ mock, setuptools, cudaPackages, - addOpenGLRunpath, + addDriverRunpath, pythonOlder, symlinkJoin, substituteAll, @@ -18,12 +18,12 @@ }: 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,7 +40,7 @@ let in buildPythonPackage rec { pname = "cupy-cuda12x"; - version = "13.1.0b"; + version = "13.3.0"; format = "setuptools"; disabled = pythonOlder "3.7"; @@ -48,8 +48,8 @@ buildPythonPackage rec { 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 index 7950c349..f08ec147 100644 --- 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 new file mode 100644 index 00000000..e0755cb8 --- /dev/null +++ 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 index c6398cd4..6ef4191e 100644 --- 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 index 4dfbd9e5..5e403607 100644 --- 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