tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 _cuda,
3 backendStdenv,
4 buildRedist,
5 cuda_cudart,
6 cudaAtLeast,
7 cudaMajorMinorVersion,
8 lib,
9 libcudla, # only for Jetson
10 patchelf,
11}:
12let
13 inherit (backendStdenv) cudaCapabilities hostRedistSystem;
14 inherit (lib.lists) optionals;
15 inherit (lib.strings) concatStringsSep optionalString;
16in
17buildRedist (
18 finalAttrs:
19 let
20 majorVersion = lib.versions.major finalAttrs.version;
21
22 tensorrtAtLeast = lib.versionAtLeast finalAttrs.version;
23 tensorrtOlder = lib.versionOlder finalAttrs.version;
24
25 # Create variables and use logical OR to allow short-circuiting.
26 tensorrtAtLeast105 = tensorrtAtLeast "10.5.0";
27 tensorrtAtLeast100 = tensorrtAtLeast105 || tensorrtAtLeast "10.0.0";
28
29 allCCNewerThan75 = lib.all (lib.flip lib.versionAtLeast "7.5") cudaCapabilities;
30 allCCNewerThan70 = allCCNewerThan75 || lib.all (lib.flip lib.versionAtLeast "7.0") cudaCapabilities;
31
32 cudaCapabilitiesJSON = builtins.toJSON cudaCapabilities;
33 in
34 {
35 redistName = "tensorrt";
36 pname = "tensorrt";
37
38 outputs = [
39 "out"
40 "bin"
41 "dev"
42 "include"
43 "lib"
44 ]
45 # From 10.14.1, TensorRT samples are distributed through the TensorRT GitHub repository.
46 ++ optionals (tensorrtOlder "10.14.1") [
47 "samples"
48 ]
49 ++ [
50 "static"
51 # "stubs" removed in postInstall
52 ];
53
54 allowFHSReferences = true;
55
56 nativeBuildInputs = [ patchelf ];
57
58 buildInputs = [
59 cuda_cudart
60 ]
61 ++ optionals libcudla.meta.available [ libcudla ];
62
63 preInstall =
64 let
65 inherit (backendStdenv.hostPlatform) parsed;
66 # x86_64-linux-gnu
67 targetString = concatStringsSep "-" [
68 parsed.cpu.name
69 parsed.kernel.name
70 parsed.abi.name
71 ];
72 in
73 # Replace symlinks to bin and lib with the actual directories from targets.
74 ''
75 for dir in bin lib; do
76 [[ -L "$dir" ]] || continue
77 nixLog "replacing symlink $PWD/$dir with $PWD/targets/${targetString}/$dir"
78 rm --verbose "$PWD/$dir"
79 mv --verbose --no-clobber "$PWD/targets/${targetString}/$dir" "$PWD/$dir"
80 done
81 unset -v dir
82 ''
83 # Remove symlinks if they exist
84 + ''
85 for dir in include samples; do
86 if [[ -L "$PWD/targets/${targetString}/$dir" ]]; then
87 nixLog "removing symlink $PWD/targets/${targetString}/$dir"
88 rm --verbose "$PWD/targets/${targetString}/$dir"
89 fi
90 done
91 unset -v dir
92
93 if [[ -d "$PWD/targets" ]]; then
94 nixLog "removing targets directory"
95 rm --recursive --verbose "$PWD/targets" || {
96 nixErrorLog "could not delete $PWD/targets: $(ls -laR "$PWD/targets")"
97 exit 1
98 }
99 fi
100 '';
101
102 autoPatchelfIgnoreMissingDeps =
103 optionals (hostRedistSystem == "linux-aarch64") [
104 "libnvdla_compiler.so"
105 ]
106 ++ optionals (tensorrtAtLeast "10.13.3") [
107 "libcuda.so.1"
108 ];
109
110 # Create a symlink for the Onnx header files in include/onnx
111 # NOTE(@connorbaker): This is shared with the tensorrt-oss package, with the `out` output swapped with `include`.
112 # When updating one, check if the other should be updated.
113 # TODO(@connorbaker): It seems like recent versions of TensorRT have separate libs for separate capabilities;
114 # we should remove libraries older than those necessary to support requested capabilities.
115 postInstall = ''
116 mkdir "''${!outputInclude:?}/include/onnx"
117 pushd "''${!outputInclude:?}/include" >/dev/null
118 nixLog "creating symlinks for Onnx header files"
119 ln -srvt "''${!outputInclude:?}/include/onnx/" NvOnnx*.h
120 popd >/dev/null
121 ''
122 # Move the python directory, which contains header files, to the include output.
123 # NOTE: Python wheels should be built from source using the TensorRT GitHub repo.
124 + ''
125 nixLog "moving python directory to include output"
126 moveToOutput python "''${!outputInclude:?}"
127
128 nixLog "remove python wheels"
129 rm --verbose "''${!outputInclude:?}"/python/*.whl
130 ''
131 + ''
132 nixLog "moving data directory to samples output"
133 moveToOutput data "''${!outputSamples:?}"
134 ''
135 # Remove the Windows library used for cross-compilation if it exists.
136 # NOTE: These are not removed for TensorRT 10.2 since the samples (and presumably others) try to load them.
137 + optionalString (lib.versions.majorMinor finalAttrs.version != "10.2") ''
138 nixLog "removing any Windows libraries"
139 for winLib in "''${!outputLib:?}/lib/"*_win*; do
140 rm --verbose "$winLib"
141 done
142 unset -v winLib
143 ''
144 # Remove the stub libraries.
145 + ''
146 nixLog "removing stub libraries"
147 rm --recursive --verbose "''${!outputLib:?}/lib/stubs" || {
148 nixErrorLog "could not delete ''${!outputLib:?}/lib/stubs"
149 exit 1
150 }
151 '';
152
153 # Tell autoPatchelf about runtime dependencies.
154 postFixup = ''
155 nixLog "patchelf-ing ''${!outputBin:?}/bin/trtexec with runtime dependencies"
156 patchelf \
157 "''${!outputBin:?}/bin/trtexec" \
158 --add-needed libnvinfer_plugin.so.${majorVersion}
159 '';
160
161 # NOTE: Like cuDNN, NVIDIA offers forward compatibility within a major releases of CUDA.
162 platformAssertions = [
163 {
164 message =
165 "tensorrt releases since 10.0.0 (found ${finalAttrs.version})"
166 + " support CUDA compute capabilities 7.0 and newer (found ${cudaCapabilitiesJSON})";
167 assertion = tensorrtAtLeast100 -> allCCNewerThan70;
168 }
169 {
170 message =
171 "tensorrt releases since 10.0.0 (found ${finalAttrs.version})"
172 + " support only CUDA compute capability 8.7 (Jetson Orin) for pre-Thor Jetson devices"
173 + " (found ${cudaCapabilitiesJSON})";
174 assertion =
175 tensorrtAtLeast100 && hostRedistSystem == "linux-aarch64" -> cudaCapabilities == [ "8.7" ];
176 }
177 {
178 message =
179 "tensorrt releases since 10.0.0 (found ${finalAttrs.version})"
180 + " support CUDA 12.4 and newer for pre-Thor Jetson devices (found ${cudaMajorMinorVersion})";
181 assertion = tensorrtAtLeast100 && hostRedistSystem == "linux-aarch64" -> cudaAtLeast "12.4";
182 }
183 {
184 message =
185 "tensorrt releases since 10.5.0 (found ${finalAttrs.version})"
186 + " support CUDA compute capabilities 7.5 and newer (found ${cudaCapabilitiesJSON})";
187 assertion = tensorrtAtLeast105 -> allCCNewerThan75;
188 }
189 ];
190
191 meta = {
192 description = "SDK that facilitates high-performance machine learning inference";
193 longDescription = ''
194 NVIDIA TensorRT is an SDK that facilitates high-performance machine learning inference. It complements training
195 frameworks such as TensorFlow, PyTorch, and MXNet. It focuses on running an already-trained network quickly and
196 efficiently on NVIDIA hardware.
197 '';
198 homepage = "https://developer.nvidia.com/tensorrt";
199 # NOTE: As of 2025-08-31, TensorRT doesn't follow the standard naming convention for URL paths that the rest of
200 # the redistributables do. As such, we need to specify downloadPage manually.
201 downloadPage = "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt";
202 changelog = "https://docs.nvidia.com/deeplearning/tensorrt/latest/getting-started/release-notes.html#release-notes";
203
204 license = _cuda.lib.licenses.tensorrt;
205 };
206 }
207)