tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 config,
3 stdenv,
4 lib,
5 fetchFromGitHub,
6 fetchpatch,
7 abseil-cpp_202407,
8 cmake,
9 cpuinfo,
10 eigen,
11 flatbuffers_23,
12 gbenchmark,
13 glibcLocales,
14 gtest,
15 howard-hinnant-date,
16 libpng,
17 nlohmann_json,
18 pkg-config,
19 python3Packages,
20 re2,
21 zlib,
22 microsoft-gsl,
23 libiconv,
24 protobuf_21,
25 pythonSupport ? true,
26 cudaSupport ? config.cudaSupport,
27 ncclSupport ? config.cudaSupport,
28 cudaPackages ? { },
29}@inputs:
30
31let
32 version = "1.22.0";
33
34 src = fetchFromGitHub {
35 owner = "microsoft";
36 repo = "onnxruntime";
37 tag = "v${version}";
38 fetchSubmodules = true;
39 hash = "sha256-fcTvMsEgO3tHOvCKCAqkO/bpZX4tcJHq9ZqpZH+uMqs=";
40 };
41
42 stdenv = throw "Use effectiveStdenv instead";
43 effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else inputs.stdenv;
44
45 cudaArchitecturesString = cudaPackages.flags.cmakeCudaArchitecturesString;
46
47 mp11 = fetchFromGitHub {
48 owner = "boostorg";
49 repo = "mp11";
50 tag = "boost-1.82.0";
51 hash = "sha256-cLPvjkf2Au+B19PJNrUkTW/VPxybi1MpPxnIl4oo4/o=";
52 };
53
54 safeint = fetchFromGitHub {
55 owner = "dcleblanc";
56 repo = "safeint";
57 tag = "3.0.28";
58 hash = "sha256-pjwjrqq6dfiVsXIhbBtbolhiysiFlFTnx5XcX77f+C0=";
59 };
60
61 pytorch_clog = effectiveStdenv.mkDerivation {
62 pname = "clog";
63 version = "3c8b153";
64 src = "${cpuinfo.src}/deps/clog";
65
66 nativeBuildInputs = [
67 cmake
68 gbenchmark
69 gtest
70 ];
71 cmakeFlags = [
72 (lib.cmakeBool "USE_SYSTEM_GOOGLEBENCHMARK" true)
73 (lib.cmakeBool "USE_SYSTEM_GOOGLETEST" true)
74 (lib.cmakeBool "USE_SYSTEM_LIBS" true)
75 # 'clog' tests set 'CXX_STANDARD 11'; this conflicts with our 'gtest'.
76 (lib.cmakeBool "CLOG_BUILD_TESTS" false)
77 ];
78 };
79
80 onnx = fetchFromGitHub {
81 owner = "onnx";
82 repo = "onnx";
83 tag = "v1.17.0";
84 hash = "sha256-9oORW0YlQ6SphqfbjcYb0dTlHc+1gzy9quH/Lj6By8Q=";
85 };
86
87 cutlass = fetchFromGitHub {
88 owner = "NVIDIA";
89 repo = "cutlass";
90 tag = "v3.5.1";
91 hash = "sha256-sTGYN+bjtEqQ7Ootr/wvx3P9f8MCDSSj3qyCWjfdLEA=";
92 };
93
94 dlpack = fetchFromGitHub {
95 owner = "dmlc";
96 repo = "dlpack";
97 tag = "v0.6";
98 hash = "sha256-YJdZ0cMtUncH5Z6TtAWBH0xtAIu2UcbjnVcCM4tfg20=";
99 };
100
101 isCudaJetson = cudaSupport && cudaPackages.flags.isJetsonBuild;
102in
103effectiveStdenv.mkDerivation rec {
104 pname = "onnxruntime";
105 inherit src version;
106
107 patches = lib.optionals cudaSupport [
108 # We apply the referenced 1064.patch ourselves to our nix dependency.
109 # FIND_PACKAGE_ARGS for CUDA was added in https://github.com/microsoft/onnxruntime/commit/87744e5 so it might be possible to delete this patch after upgrading to 1.17.0
110 ./nvcc-gsl.patch
111 ];
112
113 nativeBuildInputs = [
114 cmake
115 pkg-config
116 python3Packages.python
117 protobuf_21
118 ]
119 ++ lib.optionals pythonSupport (
120 with python3Packages;
121 [
122 pip
123 python
124 pythonOutputDistHook
125 setuptools
126 wheel
127 ]
128 )
129 ++ lib.optionals cudaSupport [
130 cudaPackages.cuda_nvcc
131 cudaPackages.cudnn-frontend
132 ]
133 ++ lib.optionals isCudaJetson [
134 cudaPackages.autoAddCudaCompatRunpath
135 ];
136
137 buildInputs = [
138 cpuinfo
139 eigen
140 glibcLocales
141 howard-hinnant-date
142 libpng
143 nlohmann_json
144 microsoft-gsl
145 pytorch_clog
146 zlib
147 ]
148 ++ lib.optionals pythonSupport (
149 with python3Packages;
150 [
151 numpy
152 pybind11
153 packaging
154 ]
155 )
156 ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [
157 libiconv
158 ]
159 ++ lib.optionals cudaSupport (
160 with cudaPackages;
161 [
162 cuda_cccl # cub/cub.cuh
163 libcublas # cublas_v2.h
164 libcurand # curand.h
165 libcusparse # cusparse.h
166 libcufft # cufft.h
167 cudnn # cudnn.h
168 cuda_cudart
169 ]
170 ++ lib.optionals (cudaSupport && ncclSupport) (
171 with cudaPackages;
172 [
173 nccl
174 ]
175 )
176 );
177
178 nativeCheckInputs = [
179 gtest
180 ]
181 ++ lib.optionals pythonSupport (
182 with python3Packages;
183 [
184 pytest
185 sympy
186 onnx
187 ]
188 );
189
190 # TODO: build server, and move .so's to lib output
191 # Python's wheel is stored in a separate dist output
192 outputs = [
193 "out"
194 "dev"
195 ]
196 ++ lib.optionals pythonSupport [ "dist" ];
197
198 enableParallelBuilding = true;
199
200 cmakeDir = "../cmake";
201
202 cmakeFlags = [
203 (lib.cmakeBool "ABSL_ENABLE_INSTALL" true)
204 (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
205 (lib.cmakeBool "FETCHCONTENT_QUIET" false)
206 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_ABSEIL_CPP" "${abseil-cpp_202407.src}")
207 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_DLPACK" "${dlpack}")
208 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_FLATBUFFERS" "${flatbuffers_23.src}")
209 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_MP11" "${mp11}")
210 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_ONNX" "${onnx}")
211 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_RE2" "${re2.src}")
212 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_SAFEINT" "${safeint}")
213 (lib.cmakeFeature "FETCHCONTENT_TRY_FIND_PACKAGE_MODE" "ALWAYS")
214 # fails to find protoc on darwin, so specify it
215 (lib.cmakeFeature "ONNX_CUSTOM_PROTOC_EXECUTABLE" "${protobuf_21}/bin/protoc")
216 (lib.cmakeBool "onnxruntime_BUILD_SHARED_LIB" true)
217 (lib.cmakeBool "onnxruntime_BUILD_UNIT_TESTS" doCheck)
218 (lib.cmakeBool "onnxruntime_USE_FULL_PROTOBUF" false)
219 (lib.cmakeBool "onnxruntime_USE_CUDA" cudaSupport)
220 (lib.cmakeBool "onnxruntime_USE_NCCL" (cudaSupport && ncclSupport))
221 (lib.cmakeBool "onnxruntime_ENABLE_LTO" (!cudaSupport || cudaPackages.cudaOlder "12.8"))
222 ]
223 ++ lib.optionals pythonSupport [
224 (lib.cmakeBool "onnxruntime_ENABLE_PYTHON" true)
225 ]
226 ++ lib.optionals cudaSupport [
227 (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${cutlass}")
228 (lib.cmakeFeature "onnxruntime_CUDNN_HOME" "${cudaPackages.cudnn}")
229 (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaArchitecturesString)
230 (lib.cmakeFeature "onnxruntime_NVCC_THREADS" "1")
231 ];
232
233 env = lib.optionalAttrs effectiveStdenv.cc.isClang {
234 NIX_CFLAGS_COMPILE = "-Wno-error";
235 };
236
237 # aarch64-linux fails cpuinfo test, because /sys/devices/system/cpu/ does not exist in the sandbox
238 doCheck = !(cudaSupport || effectiveStdenv.buildPlatform.system == "aarch64-linux");
239
240 requiredSystemFeatures = lib.optionals cudaSupport [ "big-parallel" ];
241
242 postPatch = ''
243 substituteInPlace cmake/libonnxruntime.pc.cmake.in \
244 --replace-fail '$'{prefix}/@CMAKE_INSTALL_ @CMAKE_INSTALL_
245 echo "find_package(cudnn_frontend REQUIRED)" > cmake/external/cudnn_frontend.cmake
246
247 # https://github.com/microsoft/onnxruntime/blob/c4f3742bb456a33ee9c826ce4e6939f8b84ce5b0/onnxruntime/core/platform/env.h#L249
248 substituteInPlace onnxruntime/core/platform/env.h --replace-fail \
249 "GetRuntimePath() const { return PathString(); }" \
250 "GetRuntimePath() const { return PathString(\"$out/lib/\"); }"
251 ''
252 + lib.optionalString (effectiveStdenv.hostPlatform.system == "aarch64-linux") ''
253 # https://github.com/NixOS/nixpkgs/pull/226734#issuecomment-1663028691
254 rm -v onnxruntime/test/optimizer/nhwc_transformer_test.cc
255 '';
256
257 postBuild = lib.optionalString pythonSupport ''
258 ${python3Packages.python.interpreter} ../setup.py bdist_wheel
259 '';
260
261 postInstall = ''
262 # perform parts of `tools/ci_build/github/linux/copy_strip_binary.sh`
263 install -m644 -Dt $out/include \
264 ../include/onnxruntime/core/framework/provider_options.h \
265 ../include/onnxruntime/core/providers/cpu/cpu_provider_factory.h \
266 ../include/onnxruntime/core/session/onnxruntime_*.h
267 '';
268
269 passthru = {
270 inherit cudaSupport cudaPackages; # for the python module
271 protobuf = protobuf_21;
272 tests = lib.optionalAttrs pythonSupport {
273 python = python3Packages.onnxruntime;
274 };
275 };
276
277 meta = {
278 description = "Cross-platform, high performance scoring engine for ML models";
279 longDescription = ''
280 ONNX Runtime is a performance-focused complete scoring engine
281 for Open Neural Network Exchange (ONNX) models, with an open
282 extensible architecture to continually address the latest developments
283 in AI and Deep Learning. ONNX Runtime stays up to date with the ONNX
284 standard with complete implementation of all ONNX operators, and
285 supports all ONNX releases (1.2+) with both future and backwards
286 compatibility.
287 '';
288 homepage = "https://github.com/microsoft/onnxruntime";
289 changelog = "https://github.com/microsoft/onnxruntime/releases/tag/v${version}";
290 # https://github.com/microsoft/onnxruntime/blob/master/BUILD.md#architectures
291 platforms = lib.platforms.unix;
292 license = lib.licenses.mit;
293 maintainers = with lib.maintainers; [
294 puffnfresh
295 ck3d
296 cbourjau
297 ];
298 };
299}