tjh.dev
1{
2 autoAddDriverRunpath,
3 catch2_3,
4 cmake,
5 fetchFromGitHub,
6 gitUpdater,
7 lib,
8 ninja,
9 nlohmann_json,
10 stdenv,
11 cuda_cccl ? null,
12 cuda_cudart ? null,
13 cuda_nvcc ? null,
14 cuda_nvrtc ? null,
15 cudnn ? null,
16 libcublas ? null,
17}:
18let
19 inherit (lib.lists) optionals;
20 inherit (lib.strings)
21 cmakeBool
22 cmakeFeature
23 optionalString
24 ;
25in
26
27# TODO(@connorbaker): This should be a hybrid C++/Python package.
28stdenv.mkDerivation (finalAttrs: {
29 pname = "cudnn-frontend";
30 version = "1.9.0";
31
32 src = fetchFromGitHub {
33 owner = "NVIDIA";
34 repo = "cudnn-frontend";
35 tag = "v${finalAttrs.version}";
36 hash = "sha256-Vc5jqB1XHcJEdKG0nxbWLewW2fDezRVwjUSzPDubSGE=";
37 };
38
39 patches = [
40 # https://github.com/NVIDIA/cudnn-frontend/pull/125
41 ./0001-cmake-float-out-common-python-bindings-option.patch
42 ./0002-cmake-add-config-so-headers-can-be-discovered-when-i.patch
43 ./0003-cmake-install-samples-and-tests-when-built.patch
44 ./0004-samples-fix-instances-of-maybe-uninitialized.patch
45 ];
46
47 # nlohmann_json should be the only vendored dependency.
48 postPatch = ''
49 echo "patching source to use nlohmann_json from nixpkgs"
50 rm -rf include/cudnn_frontend/thirdparty/nlohmann
51 rmdir include/cudnn_frontend/thirdparty
52 substituteInPlace include/cudnn_frontend_utils.h \
53 --replace-fail \
54 '#include "cudnn_frontend/thirdparty/nlohmann/json.hpp"' \
55 '#include <nlohmann/json.hpp>'
56 '';
57
58 # TODO: As a header-only library, we should make sure we have an `include` directory or similar which is not a
59 # superset of the `out` (`bin`) or `dev` outputs (which is what the multiple-outputs setup hook does by default).
60 outputs = [
61 "out"
62 ]
63 ++ optionals finalAttrs.doCheck [
64 "legacy_samples"
65 "samples"
66 "tests"
67 ];
68
69 nativeBuildInputs = [
70 autoAddDriverRunpath # Needed for samples because it links against CUDA::cuda_driver
71 cmake
72 cuda_nvcc
73 ninja
74 ];
75
76 buildInputs = [
77 cuda_cccl
78 cuda_cudart
79 ];
80
81 cmakeFlags = [
82 (cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
83 (cmakeFeature "FETCHCONTENT_TRY_FIND_PACKAGE_MODE" "ALWAYS")
84 (cmakeBool "CUDNN_FRONTEND_BUILD_SAMPLES" finalAttrs.doCheck)
85 (cmakeBool "CUDNN_FRONTEND_BUILD_TESTS" finalAttrs.doCheck)
86 (cmakeBool "CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS" false)
87 ];
88
89 checkInputs = [
90 cudnn
91 cuda_nvrtc
92 catch2_3
93 libcublas
94 ];
95
96 enableParallelBuilding = true;
97
98 propagatedBuildInputs = [
99 nlohmann_json
100 ];
101
102 doCheck = true;
103
104 postInstall = optionalString finalAttrs.doCheck ''
105 moveToOutput "bin/legacy_samples" "$legacy_samples"
106 moveToOutput "bin/samples" "$samples"
107 moveToOutput "bin/tests" "$tests"
108 if [[ -e "$out/bin" ]]
109 then
110 nixErrorLog "The bin directory in \$out should no longer exist."
111 exit 1
112 fi
113 '';
114
115 passthru.updateScript = gitUpdater {
116 inherit (finalAttrs) pname version;
117 rev-prefix = "v";
118 };
119
120 meta = {
121 description = "A c++ wrapper for the cudnn backend API";
122 homepage = "https://github.com/NVIDIA/cudnn-frontend";
123 license = lib.licenses.mit;
124 badPlatforms = optionals (cudnn == null) finalAttrs.meta.platforms;
125 platforms = [
126 "aarch64-linux"
127 "x86_64-linux"
128 ];
129 maintainers = with lib.maintainers; [ connorbaker ];
130 teams = [ lib.teams.cuda ];
131 };
132})