diff --git a/pkgs/development/python-modules/nvidia-cudnn-frontend/default.nix b/pkgs/development/python-modules/nvidia-cudnn-frontend/default.nix new file mode 100644 index 000000000000..4fe7bf603996 --- /dev/null +++ b/pkgs/development/python-modules/nvidia-cudnn-frontend/default.nix @@ -0,0 +1,95 @@ +{ + lib, + buildPythonPackage, + nvidia-cudnn-frontend, + + # build-system + cmake, + ninja, + pybind11, + setuptools, + + # nativeBuildInputs + cudaPackages, + + # buildInputs + dlpack, + + # propagatedBuildInputs + nlohmann_json, + + # tests + looseversion, + pytestCheckHook, + torch, +}: +buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: { + inherit (cudaPackages.cudnn-frontend) + version + src + meta + ; + + pname = "nvidia-cudnn-frontend"; + pyproject = true; + __structuredAttrs = true; + + postPatch = + cudaPackages.cudnn-frontend.postPatch + + '' + substituteInPlace pyproject.toml \ + --replace-fail '"ninja==1.11.1.1"' '"ninja"' \ + --replace-fail '"pybind11[global]>=2.13,<3"' '"pybind11"' + + sed -i '/cmake_args =/a\\f"-DCUDNN_FRONTEND_USE_SYSTEM_DLPACK=ON",' setup.py + '' + + '' + substituteInPlace python/cudnn/__init__.py \ + --replace-fail \ + 'os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/lib/libcudnn.so.*[0-9]")' \ + '"${lib.getLib cudaPackages.cudnn}/lib/libcudnn.so"' + ''; + + build-system = [ + cmake + ninja + pybind11 + setuptools + ]; + + dontUseCmakeConfigure = true; + + nativeBuildInputs = [ + cudaPackages.cuda_nvcc + ]; + + buildInputs = [ + cudaPackages.cuda_cudart + cudaPackages.cuda_nvrtc # nvrtc.h + cudaPackages.cudnn + dlpack + ]; + + propagatedBuildInputs = [ + nlohmann_json + ]; + + pythonImportsCheck = [ "cudnn" ]; + + # requires GPU + doCheck = false; + nativeCheckInputs = [ + looseversion + pytestCheckHook + torch + ]; + + enabledTestPaths = [ + "test/" + ]; + + passthru.gpuCheck = nvidia-cudnn-frontend.overridePythonAttrs { + requiredSystemFeatures = [ "cuda" ]; + doCheck = true; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 2803ff10fc33..8f166896e237 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11463,6 +11463,8 @@ self: super: with self; { nvdlib = callPackage ../development/python-modules/nvdlib { }; + nvidia-cudnn-frontend = callPackage ../development/python-modules/nvidia-cudnn-frontend { }; + nvidia-cutlass = callPackage ../development/python-modules/nvidia-cutlass { }; nvidia-cutlass-dsl = callPackage ../development/python-modules/nvidia-cutlass-dsl { };