tjh.dev
1{
2 cudaOlder,
3 cudatoolkit,
4 cudaMajorMinorVersion,
5 fetchurl,
6 lib,
7 libcublas ? null, # cuDNN uses CUDA Toolkit on old releases, where libcublas is not available.
8 patchelf,
9 zlib,
10}:
11let
12 inherit (lib)
13 lists
14 maintainers
15 meta
16 strings
17 ;
18in
19finalAttrs: prevAttrs: {
20 src = fetchurl { inherit (finalAttrs.passthru.redistribRelease) hash url; };
21
22 # Useful for inspecting why something went wrong.
23 badPlatformsConditions =
24 let
25 cudaTooOld = cudaOlder finalAttrs.passthru.featureRelease.minCudaVersion;
26 cudaTooNew =
27 (finalAttrs.passthru.featureRelease.maxCudaVersion != null)
28 && strings.versionOlder finalAttrs.passthru.featureRelease.maxCudaVersion cudaMajorMinorVersion;
29 in
30 prevAttrs.badPlatformsConditions or { }
31 // {
32 "CUDA version is too old" = cudaTooOld;
33 "CUDA version is too new" = cudaTooNew;
34 };
35
36 buildInputs =
37 prevAttrs.buildInputs or [ ]
38 ++ [ zlib ]
39 ++ lists.optionals finalAttrs.passthru.useCudatoolkitRunfile [ cudatoolkit ]
40 ++ lists.optionals (!finalAttrs.passthru.useCudatoolkitRunfile) [ (libcublas.lib or null) ];
41
42 # Tell autoPatchelf about runtime dependencies. *_infer* libraries only
43 # exist in CuDNN 8.
44 # NOTE: Versions from CUDNN releases have four components.
45 postFixup =
46 prevAttrs.postFixup or ""
47 +
48 strings.optionalString
49 (
50 strings.versionAtLeast finalAttrs.version "8.0.5.0"
51 && strings.versionOlder finalAttrs.version "9.0.0.0"
52 )
53 ''
54 ${meta.getExe patchelf} $lib/lib/libcudnn.so --add-needed libcudnn_cnn_infer.so
55 ${meta.getExe patchelf} $lib/lib/libcudnn_ops_infer.so --add-needed libcublas.so --add-needed libcublasLt.so
56 '';
57
58 passthru = prevAttrs.passthru or { } // {
59 useCudatoolkitRunfile = cudaOlder "11.3.999";
60 };
61
62 meta = prevAttrs.meta or { } // {
63 homepage = "https://developer.nvidia.com/cudnn";
64 maintainers =
65 prevAttrs.meta.maintainers or [ ]
66 ++ (with maintainers; [
67 mdaiter
68 samuela
69 connorbaker
70 ]);
71 # TODO(@connorbaker): Temporary workaround to avoid changing the derivation hash since introducing more
72 # brokenConditions would change the derivation as they're top-level and __structuredAttrs is set.
73 broken =
74 prevAttrs.meta.broken or false || (finalAttrs.passthru.useCudatoolkitRunfile && libcublas == null);
75 teams = prevAttrs.meta.teams or [ ];
76 license = {
77 shortName = "cuDNN EULA";
78 fullName = "NVIDIA cuDNN Software License Agreement (EULA)";
79 url = "https://docs.nvidia.com/deeplearning/sdk/cudnn-sla/index.html#supplement";
80 free = false;
81 redistributable = !finalAttrs.passthru.useCudatoolkitRunfile;
82 };
83 };
84}