pyrox.dev
lol
1final: prev: let
2
3 inherit (final) callPackage;
4 inherit (prev) cudatoolkit cudaVersion lib pkgs;
5
6 ### TensorRT
7
8 buildTensorRTPackage = args:
9 callPackage ./generic.nix { } args;
10
11 toUnderscore = str: lib.replaceStrings ["."] ["_"] str;
12
13 majorMinorPatch = str: lib.concatStringsSep "." (lib.take 3 (lib.splitVersion str));
14
15 tensorRTPackages = with lib; let
16 # Check whether a file is supported for our cuda version
17 isSupported = fileData: elem cudaVersion fileData.supportedCudaVersions;
18 # Return the first file that is supported. In practice there should only ever be one anyway.
19 supportedFile = files: findFirst isSupported null files;
20 # Supported versions with versions as keys and file as value
21 supportedVersions = filterAttrs (version: file: file !=null ) (mapAttrs (version: files: supportedFile files) tensorRTVersions);
22 # Compute versioned attribute name to be used in this package set
23 computeName = version: "tensorrt_${toUnderscore version}";
24 # Add all supported builds as attributes
25 allBuilds = mapAttrs' (version: file: nameValuePair (computeName version) (buildTensorRTPackage (removeAttrs file ["fileVersionCuda"]))) supportedVersions;
26 # Set the default attributes, e.g. tensorrt = tensorrt_8_4;
27 defaultBuild = { "tensorrt" = if allBuilds ? ${computeName tensorRTDefaultVersion}
28 then allBuilds.${computeName tensorRTDefaultVersion}
29 else throw "tensorrt-${tensorRTDefaultVersion} does not support your cuda version ${cudaVersion}"; };
30 in {
31 inherit buildTensorRTPackage;
32 } // allBuilds // defaultBuild;
33
34 tarballURL =
35 {fullVersion, fileVersionCuda, fileVersionCudnn ? null} :
36 "TensorRT-${fullVersion}.Linux.x86_64-gnu.cuda-${fileVersionCuda}"
37 + lib.optionalString (fileVersionCudnn != null) ".cudnn${fileVersionCudnn}"
38 + ".tar.gz";
39
40 tensorRTVersions = {
41 "8.6.0" = [
42 rec {
43 fileVersionCuda = "11.8";
44 fullVersion = "8.6.0.12";
45 sha256 = "sha256-wXMqEJPFerefoLaH8GG+Np5EnJwXeStmDzZj7Nj6e2M=";
46 tarball = tarballURL { inherit fileVersionCuda fullVersion; };
47 supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ];
48 }
49 ];
50 "8.5.3" = [
51 rec {
52 fileVersionCuda = "11.8";
53 fileVersionCudnn = "8.6";
54 fullVersion = "8.5.3.1";
55 sha256 = "sha256-BNeuOYvPTUAfGxI0DVsNrX6Z/FAB28+SE0ptuGu7YDY=";
56 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
57 supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ];
58 }
59 rec {
60 fileVersionCuda = "10.2";
61 fileVersionCudnn = "8.6";
62 fullVersion = "8.5.3.1";
63 sha256 = "sha256-WCt6yfOmFbrjqdYCj6AE2+s2uFpISwk6urP+2I0BnGQ=";
64 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
65 supportedCudaVersions = [ "10.2" ];
66 }
67 ];
68 "8.5.2" = [
69 rec {
70 fileVersionCuda = "11.8";
71 fileVersionCudnn = "8.6";
72 fullVersion = "8.5.2.2";
73 sha256 = "sha256-Ov5irNS/JETpEz01FIFNMs9YVmjGHL7lSXmDpgCdgao=";
74 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
75 supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ];
76 }
77 rec {
78 fileVersionCuda = "10.2";
79 fileVersionCudnn = "8.6";
80 fullVersion = "8.5.2.2";
81 sha256 = "sha256-UruwQShYcHLY5d81lKNG7XaoUsZr245c+PUpUN6pC5E=";
82 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
83 supportedCudaVersions = [ "10.2" ];
84 }
85 ];
86 "8.5.1" = [
87 rec {
88 fileVersionCuda = "11.8";
89 fileVersionCudnn = "8.6";
90 fullVersion = "8.5.1.7";
91 sha256 = "sha256-Ocx/B3BX0TY3lOj/UcTPIaXb7M8RFrACC6Da4PMGMHY=";
92 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
93 supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" "11.7" "11.8" ];
94 }
95 rec {
96 fileVersionCuda = "10.2";
97 fileVersionCudnn = "8.6";
98 fullVersion = "8.5.1.7";
99 sha256 = "sha256-CcFGJhw7nFdPnSYYSxcto2MHK3F84nLQlJYjdIw8dPM=";
100 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
101 supportedCudaVersions = [ "10.2" ];
102 }
103 ];
104 "8.4.0" = [
105 rec {
106 fileVersionCuda = "11.6";
107 fileVersionCudnn = "8.3";
108 fullVersion = "8.4.0.6";
109 sha256 = "sha256-DNgHHXF/G4cK2nnOWImrPXAkOcNW6Wy+8j0LRpAH/LQ=";
110 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
111 supportedCudaVersions = [ "11.0" "11.1" "11.2" "11.3" "11.4" "11.5" "11.6" ];
112 }
113 rec {
114 fileVersionCuda = "10.2";
115 fileVersionCudnn = "8.3";
116 fullVersion = "8.4.0.6";
117 sha256 = "sha256-aCzH0ZI6BrJ0v+e5Bnm7b8mNltA7NNuIa8qRKzAQv+I=";
118 tarball = tarballURL { inherit fileVersionCuda fileVersionCudnn fullVersion; };
119 supportedCudaVersions = [ "10.2" ];
120 }
121 ];
122 };
123
124 # Default attributes
125 tensorRTDefaultVersion = {
126 "10.2" = "8.4.0";
127 "11.0" = "8.4.0";
128 "11.1" = "8.4.0";
129 "11.2" = "8.4.0";
130 "11.3" = "8.4.0";
131 "11.4" = "8.4.0";
132 "11.5" = "8.4.0";
133 "11.6" = "8.4.0";
134 "11.7" = "8.5.3";
135 "11.8" = "8.5.3";
136 }.${cudaVersion} or "8.4.0";
137
138in tensorRTPackages