tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 autoAddDriverRunpath,
3 backendStdenv,
4 catch2_3,
5 cmake,
6 cuda_cccl,
7 cuda_cudart,
8 cuda_nvcc,
9 cuda_nvrtc,
10 cudaNamePrefix,
11 cudnn,
12 fetchFromGitHub,
13 gitUpdater,
14 lib,
15 libcublas,
16 ninja,
17 nlohmann_json,
18}:
19let
20 inherit (lib) licenses maintainers teams;
21 inherit (lib.lists) optionals;
22 inherit (lib.strings)
23 cmakeBool
24 cmakeFeature
25 optionalString
26 ;
27in
28# TODO(@connorbaker): This should be a hybrid C++/Python package.
29backendStdenv.mkDerivation (finalAttrs: {
30 __structuredAttrs = true;
31 strictDeps = true;
32
33 # NOTE: Depends on the CUDA package set, so use cudaNamePrefix.
34 name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
35
36 pname = "cudnn-frontend";
37 version = "1.16.0";
38
39 src = fetchFromGitHub {
40 owner = "NVIDIA";
41 repo = "cudnn-frontend";
42 tag = "v${finalAttrs.version}";
43 hash = "sha256-+8aBl9dKd2Uz50XoOr91NRyJ4OGJtzfDNNNYGQJ9b94=";
44 };
45
46 # nlohmann_json should be the only vendored dependency.
47 postPatch = ''
48 nixLog "patching source to use nlohmann_json from nixpkgs"
49 rm -rfv include/cudnn_frontend/thirdparty/nlohmann
50 rmdir -v include/cudnn_frontend/thirdparty
51 substituteInPlace include/cudnn_frontend_utils.h \
52 --replace-fail \
53 '#include "cudnn_frontend/thirdparty/nlohmann/json.hpp"' \
54 '#include <nlohmann/json.hpp>'
55 '';
56
57 # TODO: As a header-only library, we should make sure we have an `include` directory or similar which is not a
58 # superset of the `out` (`bin`) or `dev` outputs (which is what the multiple-outputs setup hook does by default).
59 outputs = [
60 "out"
61 ]
62 ++ optionals finalAttrs.doCheck [
63 "legacy_samples"
64 "samples"
65 "tests"
66 ];
67
68 nativeBuildInputs = [
69 autoAddDriverRunpath # Needed for samples because it links against CUDA::cuda_driver
70 cmake
71 cuda_nvcc
72 ninja
73 ];
74
75 buildInputs = [
76 cuda_cccl
77 cuda_cudart
78 ];
79
80 cmakeFlags = [
81 (cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
82 (cmakeFeature "FETCHCONTENT_TRY_FIND_PACKAGE_MODE" "ALWAYS")
83 (cmakeBool "CUDNN_FRONTEND_BUILD_SAMPLES" finalAttrs.doCheck)
84 (cmakeBool "CUDNN_FRONTEND_BUILD_TESTS" finalAttrs.doCheck)
85 (cmakeBool "CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS" false)
86 ];
87
88 checkInputs = [
89 cudnn
90 cuda_nvrtc
91 catch2_3
92 libcublas
93 ];
94
95 enableParallelBuilding = true;
96
97 propagatedBuildInputs = [
98 nlohmann_json
99 ];
100
101 # TODO(@connorbaker): I'm using this incorrectly to build the executables which would allow us to test functionality,
102 # rather than to indicate the checkPhase will actually run.
103 doCheck = true;
104
105 postInstall = optionalString finalAttrs.doCheck ''
106 moveToOutput "bin/legacy_samples" "$legacy_samples"
107 moveToOutput "bin/samples" "$samples"
108 moveToOutput "bin/tests" "$tests"
109 if [[ -e "$out/bin" ]]
110 then
111 nixErrorLog "The bin directory in \$out should no longer exist."
112 exit 1
113 fi
114 '';
115
116 passthru.updateScript = gitUpdater {
117 inherit (finalAttrs) pname version;
118 rev-prefix = "v";
119 };
120
121 meta = {
122 description = "A c++ wrapper for the cudnn backend API";
123 homepage = "https://github.com/NVIDIA/cudnn-frontend";
124 license = licenses.mit;
125 # Supports cuDNN 8.5.0 and newer:
126 # https://github.com/NVIDIA/cudnn-frontend/blob/11b51e9c5ad6cc71cd66cb873e34bc922d97d547/README.md?plain=1#L32
127 platforms = [
128 "aarch64-linux"
129 "x86_64-linux"
130 ];
131 maintainers = [ maintainers.connorbaker ];
132 teams = [ teams.cuda ];
133 };
134})