tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 _cuda,
3 backendStdenv,
4 cmake,
5 cuda_cudart,
6 cuda_nvcc,
7 cuda_profiler_api,
8 cudaNamePrefix,
9 fetchFromGitHub,
10 fetchzip,
11 flags,
12 lib,
13 runCommand,
14 stdenvNoCC,
15 tensorrt,
16 writeShellApplication,
17}:
18let
19 inherit (_cuda.lib) majorMinorPatch;
20 inherit (lib)
21 cmakeBool
22 cmakeFeature
23 getAttr
24 getInclude
25 licenses
26 maintainers
27 optionalString
28 replaceStrings
29 teams
30 ;
31in
32backendStdenv.mkDerivation (finalAttrs: {
33 __structuredAttrs = true;
34 strictDeps = true;
35
36 name = "${cudaNamePrefix}-${finalAttrs.pname}-${finalAttrs.version}";
37
38 pname = "tensorrt-samples";
39
40 version = majorMinorPatch tensorrt.version;
41
42 src = fetchFromGitHub (
43 {
44 owner = "NVIDIA";
45 repo = "TensorRT";
46 }
47 // getAttr finalAttrs.version {
48 "10.0.0" = {
49 tag = "v10.0.0";
50 hash = "sha256-k0FqEURPCtcPgowORHme/lhQ5SN63d0lYQvTvFXS6vw=";
51 };
52 "10.0.1" = {
53 tag = "v10.0.1";
54 hash = "sha256-lSEw0GM0eW2BHNBq/wTQA8v3aNueE3FT+k9F5nH1OgA=";
55 };
56 "10.1.0" = {
57 tag = "v10.1.0";
58 hash = "sha256-A3QwrQaI0EgRspgXEKcna/z6p7abOq3M7KDQMPQftvE=";
59 };
60 "10.2.0" = {
61 tag = "v10.2.0";
62 hash = "sha256-Euo9VD4VTpx8XJV97IMETTAx/YkPGXiNdA39Wjp3UMU=";
63 };
64 "10.3.0" = {
65 tag = "v10.3.0";
66 hash = "sha256-odSrsfOa8PlbJiPrkvWFm2hHc+ud0qSpLQanx9/M/90=";
67 };
68 "10.4.0" = {
69 tag = "v10.4.0";
70 hash = "sha256-GAu/VdHrC3UQw9okPexVItLPrRb1m3ZMpCkHNcfzRkE=";
71 };
72 "10.5.0" = {
73 tag = "v10.5.0";
74 hash = "sha256-No0JKfvi6ETXrnebMX+tAVhz7fuuCwYAp/WNUN73XzY=";
75 };
76 "10.6.0" = {
77 tag = "v10.6.0";
78 hash = "sha256-nnzicyCjVqpAonIhx3u9yNnoJkZ0XXjJ8oxQH+wfrtE=";
79 };
80 "10.7.0" = {
81 tag = "v10.7.0";
82 hash = "sha256-sbp61GverIWrHKvJV+oO9TctFTO4WUmH0oInZIwqF/s=";
83 };
84 "10.8.0" = {
85 tag = "v10.8.0";
86 hash = "sha256-SDlTZf8EQBq8vDCH3YFJCROHbf47RB5WUu4esLNpYuA=";
87 };
88 "10.9.0" = {
89 tag = "v10.9.0";
90 hash = "sha256-J8K9RjeGIem5ZxXyU+Rne8uBbul54ie6P/Y1In2mQ0g=";
91 };
92 "10.10.0" = {
93 tag = "v10.10.0";
94 hash = "sha256-/vkGmH+WKXMXsUizGfjBKRHOp5IpS236eUdCQ8tr1u8=";
95 };
96 "10.11.0" = {
97 tag = "v10.11";
98 hash = "sha256-OXI6mR2X+kF/0EO5RSBnnaGjMKD6AkuQMfl0OMzayxc=";
99 };
100 "10.12.0" = {
101 tag = "v10.12.0";
102 hash = "sha256-3pFiqDzWcMAk3GfnSOzzInddEfoGX0Fwqb+vEYr9eBw=";
103 };
104 "10.13.0" = {
105 tag = "v10.13.0";
106 hash = "sha256-hjl9fKFIE8p05/lmius2vuil6evPbNEjTT9yJyC44FI=";
107 };
108 "10.13.2" = {
109 tag = "v10.13.2";
110 hash = "sha256-1t4TyQKGTVPyPPNA3dlVDoBSHXKGcTms8AUejbvtVfw=";
111 };
112 "10.13.3" = {
113 tag = "v10.13.3";
114 hash = "sha256-d14R0UmSLT63wlmpCMi9ZvHZjottr8LJfig7EcqxLEY=";
115 };
116 "10.14.1" = {
117 tag = "v10.14";
118 hash = "sha256-pWvXpXiUriLDYHqro3HWAmO/9wbGznyUrc9qxq/t0/U=";
119 };
120 }
121 );
122
123 nativeBuildInputs = [
124 cmake
125 cuda_nvcc
126 ];
127
128 postPatch = ''
129 nixLog "patching $PWD/CMakeLists.txt to avoid manually setting CMAKE_CXX_COMPILER"
130 substituteInPlace "$PWD"/CMakeLists.txt \
131 --replace-fail \
132 'find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++)' \
133 '# find_program(CMAKE_CXX_COMPILER NAMES $ENV{CXX} g++)'
134
135 nixLog "patching $PWD/CMakeLists.txt to use find_package(CUDAToolkit) instead of find_package(CUDA)"
136 substituteInPlace "$PWD"/CMakeLists.txt \
137 --replace-fail \
138 'find_package(CUDA ''${CUDA_VERSION} REQUIRED)' \
139 'find_package(CUDAToolkit REQUIRED)'
140 ''
141 # CMakeLists.txt only started using CMAKE_CUDA_ARCHITECTURES in 10.9, and this bug was fixed by 10.12.
142 +
143 optionalString
144 (lib.versionAtLeast finalAttrs.version "10.9" && lib.versionOlder finalAttrs.version "10.12")
145 ''
146 nixLog "patching $PWD/CMakeLists.txt to fix CMake logic error"
147 substituteInPlace "$PWD"/CMakeLists.txt \
148 --replace-fail \
149 'list(APPEND CMAKE_CUDA_ARCHITECTURES SM)' \
150 'list(APPEND CMAKE_CUDA_ARCHITECTURES "''${SM}")'
151 '';
152
153 cmakeFlags = [
154 # Use tensorrt for these components; we only really want the samples.
155 (cmakeBool "BUILD_PARSERS" false)
156 (cmakeBool "BUILD_PLUGINS" false)
157 (cmakeBool "BUILD_SAMPLES" true)
158
159 # Build configuration
160 (cmakeFeature "GPU_ARCHS" (replaceStrings [ ";" ] [ " " ] flags.cmakeCudaArchitecturesString))
161 ];
162
163 buildInputs = [
164 (getInclude cuda_nvcc)
165 cuda_cudart
166 cuda_profiler_api
167 tensorrt
168 ];
169
170 passthru = import ./passthru.nix {
171 inherit
172 backendStdenv
173 fetchzip
174 finalAttrs
175 lib
176 runCommand
177 stdenvNoCC
178 writeShellApplication
179 ;
180 };
181
182 meta = {
183 description = "Open Source Software (OSS) components of NVIDIA TensorRT";
184 homepage = "https://github.com/NVIDIA/TensorRT";
185 license = licenses.asl20;
186 platforms = [
187 "aarch64-linux"
188 "x86_64-linux"
189 ];
190 teams = [ teams.cuda ];
191 maintainers = with maintainers; [ connorbaker ];
192 };
193})