Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185{ lib, buildPythonPackage, cudaPackages, fetchFromGitHub,
# build-system cython, setuptools,
# nativeBuildInputs autoAddDriverRunpath, writableTmpDirAsHomeHook,
# dependencies numpy, cuda-pathfinder,}:
buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: { pname = "cupy"; version = "14.2.0"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "cupy"; repo = "cupy"; tag = "v${finalAttrs.version}"; fetchSubmodules = true; hash = "sha256-EYje+FnIVIsd+rVrCkwXXyLp6jnPTTFpuRRri88XK0k="; };
patches = [ # Let cupy_setup_build.py find static libraries (.a) from buildInputs. # By default, it looks for them in `$CUDA_PATH/lib{64,}`. ./link-static-libraries.patch ];
postPatch = # Relax cython constraint '' substituteInPlace pyproject.toml \ --replace-fail \ '"Cython>=3.2,<3.3",' \ '"Cython>=3.2",' '' # Inject absolute path to nvcc instead of relying on runtime discovery heuristics + '' substituteInPlace cupy/_environment.py \ --replace-fail \ "nvcc_path = os.environ.get('NVCC', None)" \ "return '${lib.getExe cudaPackages.cuda_nvcc}'" '' # Inject absolute path to CUDA libraries which get disovered at runtime + '' substituteInPlace cupy/_core/core.pyx \ --replace-fail \ "_cuda_include_dir = find_nvidia_header_directory('cudart')" \ "_cuda_include_dir = '${lib.getInclude cudaPackages.cuda_cudart}/include'"
substituteInPlace cupy/cuda/__init__.py \ --replace-fail \ "from cuda import pathfinder" \ "from ctypes import CDLL" \ --replace-fail \ 'pathfinder.load_nvidia_dynamic_lib("cufft")' \ 'CDLL("${lib.getLib cudaPackages.libcufft}/lib/libcufft.so")'
substituteInPlace cupy/cuda/cufft.pyx \ --replace-fail \ "from cuda import pathfinder" \ "from ctypes import CDLL" \ --replace-fail \ "loaded_dl = pathfinder.load_nvidia_dynamic_lib('cufft')" \ "loaded_dl = CDLL('${lib.getLib cudaPackages.libcufft}/lib/libcufft.so')" \ --replace-fail \ "handle = loaded_dl._handle_uint" \ "handle = loaded_dl._handle"
substituteInPlace cupy_backends/cuda/api/_runtime_softlink.pxi \ --replace-fail \ "from cuda import pathfinder" \ "from ctypes import CDLL" \ --replace-fail \ "loaded_dl = pathfinder.load_nvidia_dynamic_lib('cudart')" \ "loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_cudart}/lib/libcudart.so')" \ --replace-fail \ "handle = loaded_dl._handle_uint" \ "handle = loaded_dl._handle"
substituteInPlace cupy_backends/cuda/libs/_cnvrtc.pxi \ --replace-fail \ "from cuda import pathfinder" \ "from ctypes import CDLL" \ --replace-fail \ "loaded_dl = pathfinder.load_nvidia_dynamic_lib('nvrtc')" \ "loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so')" \ --replace-fail \ "handle = loaded_dl._handle_uint" \ "handle = loaded_dl._handle"
substituteInPlace cupy/cuda/compiler.py \ --replace-fail \ "cudadevrt = get_cuda_path()" \ "return '${lib.getOutput "static" cudaPackages.cuda_cudart}/lib/libcudadevrt.a'" '';
env = { LDFLAGS = toString [ # Fake libcuda.so (the real one is deployed impurely) "-L${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" ]; };
# See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Setting 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" '';
enableParallelBuilding = true;
build-system = [ cython setuptools ];
nativeBuildInputs = [ autoAddDriverRunpath cudaPackages.cuda_nvcc
writableTmpDirAsHomeHook # Needed for pythonImportsCheck ];
buildInputs = with cudaPackages; [ cuda_cudart # cuda.h, cuda_runtime.h, libcudart_static.a cuda_cuxxfilt # nv_decode.h cuda_nvrtc # nvrtc.h nvtx # nvToolsExt.h cuda_profiler_api # cuda_profiler_api.h libcublas # cublas_v2.h libcufft # cufft.h libcurand # curand.h libcusolver # cusolverDn.h libcusparse # cusparse.h libcusparse_lt # cusparseLt.h libcutensor # cutensor.h nccl # nccl.h
(lib.getOutput "static" cuda_cuxxfilt) # libcufilt.a ];
dependencies = [ cuda-pathfinder numpy ];
pythonImportsCheck = [ "cupy" "cupy_backends" "cupyx" ];
# Won't work with the GPU, whose drivers won't be accessible from the build sandbox doCheck = false;
meta = { description = "NumPy-compatible matrix library accelerated by CUDA"; homepage = "https://cupy.chainer.org/"; changelog = "https://github.com/cupy/cupy/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.mit; platforms = [ "aarch64-linux" "x86_64-linux" ]; maintainers = with lib.maintainers; [ GaetanLepage ]; teams = [ lib.teams.cuda ]; };})