From 1dd5702a1355c2c143ab8faa96aa0e323237583b Mon Sep 17 00:00:00 2001 From: Jake Arkinstall <65358059+jake-arkinstall@users.noreply.github.com> Date: Mon, 24 Jun 2024 15:04:17 +0100 Subject: [PATCH] Added basic nix support via a flake This commit utilises some cuda and python packages from nixpkgs to fetch dependencies, and also provides missing dependencies and upgrades to several dependencies. It allows one to run `nix build .#pytket-cutensornet` (or `nix build` for short) to build the tket extension. Tests can't be run within the normal `nix flake check` context as the GPU drivers aren't available within the sandbox, so tests can be run with `nix run .#tests` instead. `nix develop` within the source root can be used to enter a shell with pytket's cutensornet backend available. If and when this commit is merged into main, users will be able to run: `nix develop github:cqcl/pytket-cutensornet` and have the environment fully loaded for their use. Cuda support is verified on nixos. Verification on non-nixos nix installations will be required for general use. Integration with Cachix will be extremely beneficial due to the size and duration of the builds. --- .gitignore | 6 +- flake.lock | 115 +++++++++++++++++++++++++ flake.nix | 52 +++++++++++ nix-support/cuda-bundle.nix | 37 ++++++++ nix-support/cupy.nix | 134 +++++++++++++++++++++++++++++ nix-support/cupy.patch | 36 ++++++++ nix-support/cuquantum.nix | 50 +++++++++++ nix-support/cusparselt.nix | 34 ++++++++ nix-support/cutensor.nix | 37 ++++++++ nix-support/pycuquantum.nix | 45 ++++++++++ nix-support/pycuquantum.patch | 62 +++++++++++++ nix-support/pytket-cutensornet.nix | 66 ++++++++++++++ 12 files changed, 673 insertions(+), 1 deletion(-) create mode 100644 flake.lock create mode 100644 flake.nix create mode 100644 nix-support/cuda-bundle.nix create mode 100644 nix-support/cupy.nix create mode 100644 nix-support/cupy.patch create mode 100644 nix-support/cuquantum.nix create mode 100644 nix-support/cusparselt.nix create mode 100644 nix-support/cutensor.nix create mode 100644 nix-support/pycuquantum.nix create mode 100644 nix-support/pycuquantum.patch create mode 100644 nix-support/pytket-cutensornet.nix diff --git a/.gitignore b/.gitignore index f7a86986..3b25b435 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,8 @@ dist obj docs/extensions .ipynb_checkpoints -pytket/extensions/cutensornet/_metadata.py \ No newline at end of file +pytket/extensions/cutensornet/_metadata.py +# nix artifacts +result +.direnv +.envrc diff --git a/flake.lock b/flake.lock new file mode 100644 index 00000000..d3959756 --- /dev/null +++ b/flake.lock @@ -0,0 +1,115 @@ +{ + "nodes": { + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1710146030, + "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "flake-utils_2": { + "inputs": { + "systems": "systems_2" + }, + "locked": { + "lastModified": 1710146030, + "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1719235230, + "narHash": "sha256-o2mvJJxQ0dE7YHJsWjO8nMfWCvZQHAC9DAtiMc3GSBs=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "6acfb07b7ff392f1cdfe2dfb96f21cb919ca359d", + "type": "github" + }, + "original": { + "owner": "NixOS", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "tket": "tket" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + }, + "systems_2": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + }, + "tket": { + "inputs": { + "flake-utils": "flake-utils_2", + "nixpkgs": [ + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1718711031, + "narHash": "sha256-sAYbO/qXZ61Y3Yk2K1VQU+ER+dGOgIZeeNM9FPe27iI=", + "owner": "CQCL", + "repo": "tket", + "rev": "d5451418f257e0aee61244aed56714b894c0ef09", + "type": "github" + }, + "original": { + "owner": "CQCL", + "repo": "tket", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 00000000..aba8a503 --- /dev/null +++ b/flake.nix @@ -0,0 +1,52 @@ +{ + description = "Pytket Cutensornet Extension"; + nixConfig.extra-substituters = "https://tket.cachix.org https://cache.nixos.org https://cuda-maintainers.cachix.org"; + nixConfig.trusted-public-keys = '' + tket.cachix.org-1:ACdm5Zg19qPL0PpvUwTPPiIx8SEUy+D/uqa9vKJFwh0= + cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY= + cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E= + ''; + inputs.nixpkgs.url = "github:NixOS/nixpkgs"; + inputs.flake-utils.url = "github:numtide/flake-utils"; + inputs.tket.url = "github:CQCL/tket"; + inputs.tket.inputs.nixpkgs.follows = "nixpkgs"; + outputs = { self, nixpkgs, flake-utils, tket }: + flake-utils.lib.eachDefaultSystem (system: + let + pkgs = import nixpkgs { + inherit system; + config = { + allowUnfree = true; + cudaSupport = true; + }; + overlays = [ + (self: super: { + inherit (tket.packages."${system}") tket pytket; + }) + (self: super: { + cuda-bundle = pkgs.callPackage ./nix-support/cuda-bundle.nix {}; + }) + (self: super: { + cupy' = pkgs.python3Packages.callPackage ./nix-support/cupy.nix {}; + }) + (self: super: { + pycuquantum = pkgs.python3Packages.callPackage ./nix-support/pycuquantum.nix {}; + }) + (import ./nix-support/pytket-cutensornet.nix) + ]; + }; + in { + packages = { + default = pkgs.pytket-cutensornet; + cupy = pkgs.cupy'; + pytket-cutensornet = pkgs.pytket-cutensornet; + tests = pkgs.run-pytket-cutensornet-tests; + }; + devShells = { + default = pkgs.mkShell { buildInputs = [ pkgs.pytket-cutensornet ]; }; + }; + 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 new file mode 100644 index 00000000..ad2ad430 --- /dev/null +++ b/nix-support/cuda-bundle.nix @@ -0,0 +1,37 @@ +{ + symlinkJoin, + cudaPackages, +}: +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 + cuda_nvprof + cuda_nvrtc + cuda_nvtx + cuda_profiler_api + libcublas + libcufft + libcurand + libcusolver + libcusparse + cusparselt' + cusparselt'.dev + cutensor' + cutensor'.dev + cuquantum' + ]; + postBuild = '' + ln -sfn lib $out/lib64 + ''; +} + diff --git a/nix-support/cupy.nix b/nix-support/cupy.nix new file mode 100644 index 00000000..932a7cb5 --- /dev/null +++ b/nix-support/cupy.nix @@ -0,0 +1,134 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + cython_0, + fastrlock, + numpy, + wheel, + pytestCheckHook, + mock, + setuptools, + cudaPackages, + addOpenGLRunpath, + pythonOlder, + symlinkJoin, + substituteAll, + cuda-bundle +}: + +let + inherit (cudaPackages) cudnn nccl; + cupy-cccl-src = fetchFromGitHub{ + owner = "cupy"; + repo = "cccl"; + rev = "79ed0e96e35112d171e43f13fa7f324eff7f3de0"; + hash = "sha256-8YlIp7xOQtB7dKhCL7SIX+cKOvP4V/XPx6EC+Ct4vSc="; + }; + dlpack-src = fetchFromGitHub{ + owner = "dmlc"; + repo = "dlpack"; + rev = "365b823cedb281cd0240ca601aba9b78771f91a3"; + hash = "sha256-IcfCoz3PfDdRetikc2MZM1sJFOyRgKonWMk21HPbrso="; + }; + jitify-src = fetchFromGitHub{ + owner = "NVIDIA"; + repo = "jitify"; + rev = "1a0ca0e837405506f3b8f7883bacb71c20d86d96"; + sha256 = "sha256-m+ZrYbmA/FuKZVIUiLYX5GR79jheu5CNHtli/X+ux0U="; + }; +in +buildPythonPackage rec { + pname = "cupy-cuda12x"; + version = "13.1.0b"; + format = "setuptools"; + + disabled = pythonOlder "3.7"; + + src = fetchFromGitHub { + owner = "cupy"; + repo = "cupy"; + rev = "6730353d611f4fd5f6a7494e90c7639a35245337"; # cusparse_lt 0.6.1 support + hash = sha256:h6Q6abnq5zPY6mx8dA5l1kXfKjisEOwEqP0rGGu9q0M=; + }; + patches = [ + (substituteAll { + src = ./cupy.patch; + cuda_bundle = cuda-bundle.out; + }) + ]; + + unpackPhase = '' + for f in $src/*; do + if [ "$f" != "$src/third_party" ]; then + cp -r $f .; + fi + done + mkdir third_party; + ln -s ${dlpack-src} third_party/dlpack; + ln -s ${jitify-src} third_party/jitify; + ln -s ${cupy-cccl-src} third_party/cccl; + chmod -R 0775 .; + ''; + + # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both + # CUPY_NUM_BUILD_JOBS and CUPY_NUM_NVCC_THREADS to NIX_BUILD_CORES results in + # a small amount of thrashing but it turns out there are a large number of + # very short builds and a few extremely long ones, so setting both ends up + # working nicely in practice. + preConfigure = '' + export CUPY_NUM_BUILD_JOBS="$NIX_BUILD_CORES" + export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES" + ''; + + nativeBuildInputs = [ + cuda-bundle + setuptools + wheel + addOpenGLRunpath + cython_0 + ]; + + buildInputs = [ + cuda-bundle + cudnn + nccl + ]; + + NVCC="${cuda-bundle}/bin/nvcc"; + CUDA_PATH="${cuda-bundle}"; + LDFLAGS="-L${cuda-bundle}/lib"; + LD_LIBRARY_PATH="${cuda-bundle}/lib"; + + propagatedBuildInputs = [ + fastrlock + numpy + ]; + + nativeCheckInputs = [ + pytestCheckHook + mock + ]; + + # Won't work with the GPU, whose drivers won't be accessible from the build + # sandbox + doCheck = false; + + postFixup = '' + find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do + addOpenGLRunpath "$lib" + patchelf --set-rpath ${cuda-bundle}/lib $lib; + done + ''; + + enableParallelBuilding = true; + + meta = with lib; { + description = "A NumPy-compatible matrix library accelerated by CUDA"; + homepage = "https://cupy.chainer.org/"; + changelog = "https://github.com/cupy/cupy/releases/tag/v${version}"; + license = licenses.mit; + platforms = [ "x86_64-linux" ]; + maintainers = with maintainers; [ hyphon81 ]; + }; +} diff --git a/nix-support/cupy.patch b/nix-support/cupy.patch new file mode 100644 index 00000000..7950c349 --- /dev/null +++ b/nix-support/cupy.patch @@ -0,0 +1,36 @@ +diff --git a/cupy_backends/cuda/libs/_cnvrtc.pxi b/cupy_backends/cuda/libs/_cnvrtc.pxi +index 9f02b5522..3e0ef9da9 100644 +--- a/cupy_backends/cuda/libs/_cnvrtc.pxi ++++ b/cupy_backends/cuda/libs/_cnvrtc.pxi +@@ -114,30 +114,5 @@ cdef void _initialize() except *: + cdef SoftLink _get_softlink(): + cdef int runtime_version + cdef str prefix = 'nvrtc' +- cdef object libname = None +- +- if CUPY_CUDA_VERSION != 0: +- runtime_version = runtime._getCUDAMajorVersion() +- if runtime_version == 11: +- # CUDA 11.x (11.2+) +- if _sys.platform == 'linux': +- libname = 'libnvrtc.so.11.2' +- else: +- libname = 'nvrtc64_112_0.dll' +- elif runtime_version == 12: +- # CUDA 12.x +- if _sys.platform == 'linux': +- libname = 'libnvrtc.so.12' +- else: +- libname = 'nvrtc64_120_0.dll' +- elif CUPY_HIP_VERSION != 0: +- runtime_version = runtime.runtimeGetVersion() +- prefix = 'hiprtc' +- if runtime_version < 5_00_00000: +- # ROCm 4.x +- libname = 'libamdhip64.so.4' +- elif runtime_version < 6_00_00000: +- # ROCm 5.x +- libname = 'libamdhip64.so.5' +- ++ cdef object libname = '@cuda_bundle@/lib/libnvrtc.so' + return SoftLink(libname, prefix, mandatory=True) diff --git a/nix-support/cuquantum.nix b/nix-support/cuquantum.nix new file mode 100644 index 00000000..15389d08 --- /dev/null +++ b/nix-support/cuquantum.nix @@ -0,0 +1,50 @@ +{ + pkgs, + system, + cudaPackages, + fetchzip, + stdenv, + patchelf, + #cuda-bundle, +}: +let + version = "24.03.0.4"; + package_map = { + x86_64-linux = { + "11" = fetchzip { + url = "https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-x86_64/cuquantum-linux-x86_64-${version}_cuda11-archive.tar.xz"; + sha256 = sha256:tvUEKdUp3dluY7uzm6NJBZxcjJsVhMtwJmV061vR0XU=; + }; + "12" = fetchzip { + url = "https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-x86_64/cuquantum-linux-x86_64-${version}_cuda12-archive.tar.xz"; + sha256 = sha256:ws9dPv4Tkdu9XSE4zI3WdFccQuhiycrEchduzPNBiCc=; + }; + }; + aarch64-linux = { + "11" = fetchzip { + url = "https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-sbsa/cuquantum-linux-sbsa-${version}_cuda11-archive.tar.xz"; + sha256 = sha256:Wc0HgfnDgTpyDKMQpnJdJ95FkBe+7XAqQ48XkkrWiKU=; + }; + "12" = fetchzip { + url = "https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-sbsa/cuquantum-linux-sbsa-${version}_cuda12-archive.tar.xz"; + sha256 = sha256:ZnjnMja0565Mi9GE6IE3uDXXY1jisz/VS1Yl7nZSaVw=; + }; + }; + }; +in + stdenv.mkDerivation { + pname = "cuquantum"; + inherit version; + src = package_map.${system}.${cudaPackages.cudaMajorVersion}; + nativeBuildInputs = [ pkgs.patchelf ]; + installPhase = '' + mkdir -p $out + cp -r $src/* $out + chmod -R 0755 $out + ''; + #fixupPhase = '' + # for f in $out/lib/*.so; do + # patchelf --set-rpath ${cuda-bundle}/lib $f; + # done; + #''; + } diff --git a/nix-support/cusparselt.nix b/nix-support/cusparselt.nix new file mode 100644 index 00000000..b39c0899 --- /dev/null +++ b/nix-support/cusparselt.nix @@ -0,0 +1,34 @@ +{ pkgs, cudaPackages, fetchzip }: +let + package_map = { + x86_64-linux = fetchzip { + url = "https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.6.1.0-archive.tar.xz"; + sha256 = sha256:+CuYbJ2Sac3kmzdrJQLvdl2oGehx3b6kQqcyHksYxyU=; + }; + aarch64-linux = fetchzip { + url = "https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-spsa/libcusparse_lt-linux-spsa-0.6.1.0-archive.tar.xz"; + sha256 = ""; #TODO + }; + }; +in pkgs.stdenv.mkDerivation { + pname = "libcusparse_lt"; + version = "0.6.1.0"; + outputs = [ "out" "lib" "dev" "static" ]; + src = package_map.${pkgs.system}; + installPhase = '' + mkdir -p $lib/lib; + cp -r $src/lib/*.so $lib/lib; + cp -r $src/lib/*.so.* $lib/lib; + + mkdir -p $dev/include; + cp -r $src/include $dev/include; + + mkdir -p $static/lib; + cp -r $src/lib/libcusparseLt_static.a $static/lib; + + mkdir -p $out; + cp -r $src/include $out/include; + cp -r $src/lib $out/lib; + cp $src/LICENSE $out/LICENSE; + ''; +} diff --git a/nix-support/cutensor.nix b/nix-support/cutensor.nix new file mode 100644 index 00000000..75245762 --- /dev/null +++ b/nix-support/cutensor.nix @@ -0,0 +1,37 @@ +{ pkgs, cudaPackages, fetchzip }: +let + package_map = { + x86_64-linux = fetchzip { + url = "https://developer.download.nvidia.com/compute/cutensor/redist/libcutensor/linux-x86_64/libcutensor-linux-x86_64-2.0.1.2-archive.tar.xz"; + sha256 = sha256:OPn5577DBwzcGdSovv5tMv+x3wW2QgF5XVmRgm5y148=; + }; + aarch64-linux = fetchzip { + url = "https://developer.download.nvidia.com/compute/cutensor/redist/libcutensor/linux-sbsa/libcutensor-linux-sbsa-2.0.1.2-archive.tar.xz"; + sha256 = ""; #TODO + }; + }; +in pkgs.stdenv.mkDerivation { + pname = "libcutensor"; + version = "2.0.1.2"; + outputs = [ "out" "lib" "dev" "static" ]; + src = package_map.${pkgs.system}; + installPhase = let + v = cudaPackages.cudaMajorVersion; + in '' + mkdir -p $lib/lib; + cp -r $src/lib/${v}/*.so $lib/lib; + cp -r $src/lib/${v}/*.so.* $lib/lib; + + mkdir -p $dev/include; + cp -r $src/include $dev/include; + + mkdir -p $static/lib; + cp -r $src/lib/${v}/libcutensorMg_static.a $static/lib; + cp -r $src/lib/${v}/libcutensor_static.a $static/lib; + + mkdir -p $out; + cp -r $src/include $out/include; + cp -r $src/lib/${v} $out/lib; + cp $src/LICENSE $out/LICENSE; + ''; +} diff --git a/nix-support/pycuquantum.nix b/nix-support/pycuquantum.nix new file mode 100644 index 00000000..a2ee986e --- /dev/null +++ b/nix-support/pycuquantum.nix @@ -0,0 +1,45 @@ +{ + python3Packages, + buildPythonPackage, + fetchFromGitHub, + cupy', + wheel, + cython_0, + setuptools, + cuda-bundle, + substituteAll +}: +buildPythonPackage{ + pname = "cuquantum"; + version = "24.03.0"; + pyproject = true; + src = fetchFromGitHub { + owner = "NVIDIA"; + repo = "cuQuantum"; + rev = "v24.03.0"; + sha256 = sha256:1GPotBCNEsriU+yrWQpIVWKh/wEK8EPWYjn3jF4w8T0=; + }; + patches = [ + (substituteAll { + src = ./pycuquantum.patch; + cuda_bundle = cuda-bundle.out; + }) + ]; + build-system = [ + setuptools + wheel + cython_0 + ]; + dependencies = [ + cupy' + ]; + dontCheckRuntimeDeps = true; + CUDA_PATH="${cuda-bundle}"; + CUQUANTUM_ROOT="${cuda-bundle}"; + CUTENSOR_ROOT="${cuda-bundle}"; + + unpackPhase = '' + cp -r $src/python/* . + chmod -R 0755 . + ''; +} diff --git a/nix-support/pycuquantum.patch b/nix-support/pycuquantum.patch new file mode 100644 index 00000000..dff46380 --- /dev/null +++ b/nix-support/pycuquantum.patch @@ -0,0 +1,62 @@ +diff --git a/builder/pep517.py b/builder/pep517.py +index ce701e4..8b99806 100644 +--- a/builder/pep517.py ++++ b/builder/pep517.py +@@ -30,11 +30,12 @@ def get_requires_for_build_wheel(config_settings=None): + # set up version constraints: note that CalVer like 22.03 is normalized to + # 22.3 by setuptools, so we must follow the same practice in the constraints; + # also, we don't need the patch number here +- cuqnt_require = [f'custatevec-cu{utils.cuda_major_ver}~=1.6', # ">=1.6.0,<2" +- f'cutensornet-cu{utils.cuda_major_ver}~=2.4', # ">=2.4.0,<3" +- ] ++ #cuqnt_require = [f'custatevec-cu{utils.cuda_major_ver}~=1.6', # ">=1.6.0,<2" ++ # f'cutensornet-cu{utils.cuda_major_ver}~=2.4', # ">=2.4.0,<3" ++ # ] + +- return _build_meta.get_requires_for_build_wheel(config_settings) + cuqnt_require ++ #return _build_meta.get_requires_for_build_wheel(config_settings) + cuqnt_require ++ return _build_meta.get_requires_for_build_wheel(config_settings) + + + # Note: We have never promised to support sdist (CUQNT-514). We really cannot +diff --git a/cuquantum/custatevec/_internal/custatevec_linux.pyx b/cuquantum/custatevec/_internal/custatevec_linux.pyx +index 8b918ca..afc6d87 100644 +--- a/cuquantum/custatevec/_internal/custatevec_linux.pyx ++++ b/cuquantum/custatevec/_internal/custatevec_linux.pyx +@@ -111,7 +111,7 @@ cdef void* __custatevecSubSVMigratorMigrate = NULL + + cdef void* load_library() except* nogil: + cdef void* handle +- handle = dlopen("libcustatevec.so.1", RTLD_NOW | RTLD_GLOBAL) ++ handle = dlopen("@cuda_bundle@/lib/libcustatevec.so", RTLD_NOW | RTLD_GLOBAL) + if handle == NULL: + with gil: + err_msg = dlerror() +diff --git a/cuquantum/cutensornet/_internal/cutensornet_linux.pyx b/cuquantum/cutensornet/_internal/cutensornet_linux.pyx +index b63a366..1143833 100644 +--- a/cuquantum/cutensornet/_internal/cutensornet_linux.pyx ++++ b/cuquantum/cutensornet/_internal/cutensornet_linux.pyx +@@ -152,7 +152,7 @@ cdef void* __cutensornetSamplerGetInfo = NULL + + cdef void* load_library() except* nogil: + cdef void* handle +- handle = dlopen("libcutensornet.so.2", RTLD_NOW | RTLD_GLOBAL) ++ handle = dlopen("@cuda_bundle@/lib/libcutensornet.so", RTLD_NOW | RTLD_GLOBAL) + if handle == NULL: + with gil: + err_msg = dlerror() +diff --git a/setup.py b/setup.py +index 47a892f..11f59db 100644 +--- a/setup.py ++++ b/setup.py +@@ -35,8 +35,8 @@ with open(os.path.join(source_root, "tests/requirements.txt")) as f: + install_requires = [ + 'numpy~=1.21', # ">=1.21,<2" + # 'torch', # <-- PyTorch is optional; also, the PyPI version does not support GPU... +- f'custatevec-cu{utils.cuda_major_ver}~=1.6', # ">=1.6.0,<2" +- f'cutensornet-cu{utils.cuda_major_ver}~=2.4', # ">=2.4.0,<3" ++ #f'custatevec-cu{utils.cuda_major_ver}~=1.6', # ">=1.6.0,<2" ++ #f'cutensornet-cu{utils.cuda_major_ver}~=2.4', # ">=2.4.0,<3" + ] + if utils.cuda_major_ver == '11': + install_requires.append('cupy-cuda11x>=13.0') # no ambiguity diff --git a/nix-support/pytket-cutensornet.nix b/nix-support/pytket-cutensornet.nix new file mode 100644 index 00000000..d95840e7 --- /dev/null +++ b/nix-support/pytket-cutensornet.nix @@ -0,0 +1,66 @@ +self: super: +let + metadata = builtins.readFile ../_metadata.py; + versions = + builtins.match ''.*_version__ *= *["']([^"']+)["'].*'' metadata; + version = if builtins.length versions > 0 then + builtins.elemAt versions 0 + else + builtins.trace "Warning: unable to find version. Defaulting to 0.0.0" "0.0.0"; +in { + pytket-cutensornet = super.python3Packages.buildPythonPackage { + pname = "pytket-cutensornet"; + version = version; + src = super.stdenv.mkDerivation{ + name = "pytket-cutensornet-sources"; + phases = [ "installPhase" ]; + installPhase = '' + mkdir -p $out; + cp -r ${../pytket} $out/pytket; + + cp ${../setup.py} $out/setup.py; + cp ${../README.md} $out/README.md; # required for setup's long description + cp ${../pytest.ini} $out/pytest.ini; + cp ${../_metadata.py} $out/_metadata.py; + + # on nix versions of scipy and ipython, stubs are missing. + # adjust mypy.ini to ignore these errors. + ( + cat ${../mypy.ini}; + cat <> $out/mypy.ini; + ''; + }; + propagatedBuildInputs = [ super.pytket super.pycuquantum ]; + + doCheck = true; + checkInputs = with super.python3Packages; [ 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 + ''; + }; + + run-pytket-cutensornet-tests = let + test-env = super.python3.withPackages(ps: with ps; [ + self.pytket-cutensornet + pytest + pytest-lazy-fixture + ]); + in super.writeShellScriptBin "run-pytket-cutensornet-tests" '' + HOME=$(mktemp -d); + export HOME; + echo "---------------------------"; + env; + echo "---------------------------"; + ${test-env}/bin/pytest -s ${../tests}; + ''; +}