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

@@ -9,4 +9,8 @@ dist
 obj
 docs/extensions
 .ipynb_checkpoints
-pytket/extensions/cutensornet/_metadata.py
+pytket/extensions/cutensornet/_metadata.py
+# nix artifacts
+result
+.direnv
+.envrc

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
+        };
+      });
+}

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 # <nv/target>
+    cuda_cccl.dev
+    cuda_cudart
+    cuda_nvcc.dev # <crt/host_defines.h>
+    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
+  '';
+}
+

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

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)