pkgs/development/libraries/onnxruntime/default.nix

   1 { config
   2 , stdenv
   3 , lib
   4 , fetchFromGitHub
   5 , Foundation
   6 , abseil-cpp_202401
   7 , cmake
   8 , cpuinfo
   9 , eigen
  10 , flatbuffers_23
  11 , gbenchmark
  12 , glibcLocales
  13 , gtest
  14 , libpng
  15 , nlohmann_json
  16 , nsync
  17 , pkg-config
  18 , python3Packages
  19 , re2
  20 , zlib
  21 , microsoft-gsl
  22 , libiconv
  23 , protobuf_21
  24 , pythonSupport ? true
  25 , cudaSupport ? config.cudaSupport
  26 , ncclSupport ? config.cudaSupport
  27 , cudaPackages ? {}
  28 }@inputs:
  29
  30
  31 let
  32   version = "1.18.1";
  33
  34   abseil-cpp = abseil-cpp_202401;
  35
  36   stdenv = throw "Use effectiveStdenv instead";
  37   effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else inputs.stdenv;
  38
  39   cudaArchitecturesString = cudaPackages.flags.cmakeCudaArchitecturesString;
  40
  41   howard-hinnant-date = fetchFromGitHub {
  42     owner = "HowardHinnant";
  43     repo = "date";
  44     rev = "v3.0.1";
  45     sha256 = "sha256-ZSjeJKAcT7mPym/4ViDvIR9nFMQEBCSUtPEuMO27Z+I=";
  46   };
  47
  48   mp11 = fetchFromGitHub {
  49     owner = "boostorg";
  50     repo = "mp11";
  51     rev = "boost-1.82.0";
  52     hash = "sha256-cLPvjkf2Au+B19PJNrUkTW/VPxybi1MpPxnIl4oo4/o=";
  53   };
  54
  55   safeint = fetchFromGitHub {
  56     owner = "dcleblanc";
  57     repo = "safeint";
  58     rev = "ff15c6ada150a5018c5ef2172401cb4529eac9c0";
  59     hash = "sha256-PK1ce4C0uCR4TzLFg+elZdSk5DdPCRhhwT3LvEwWnPU=";
  60   };
  61
  62   pytorch_clog = effectiveStdenv.mkDerivation {
  63     pname = "clog";
  64     version = "3c8b153";
  65     src = "${cpuinfo.src}/deps/clog";
  66
  67     nativeBuildInputs = [ cmake gbenchmark gtest ];
  68     cmakeFlags = [
  69       "-DUSE_SYSTEM_GOOGLEBENCHMARK=ON"
  70       "-DUSE_SYSTEM_GOOGLETEST=ON"
  71       "-DUSE_SYSTEM_LIBS=ON"
  72       # 'clog' tests set 'CXX_STANDARD 11'; this conflicts with our 'gtest'.
  73       "-DCLOG_BUILD_TESTS=OFF"
  74     ];
  75   };
  76
  77   onnx = fetchFromGitHub {
  78     owner = "onnx";
  79     repo = "onnx";
  80     rev = "refs/tags/v1.16.1";
  81     hash = "sha256-I1wwfn91hdH3jORIKny0Xc73qW2P04MjkVCgcaNnQUE=";
  82   };
  83
  84    cutlass = fetchFromGitHub {
  85     owner = "NVIDIA";
  86     repo = "cutlass";
  87     rev = "v3.1.0";
  88     hash = "sha256-mpaiCxiYR1WaSSkcEPTzvcREenJWklD+HRdTT5/pD54=";
  89  };
  90 in
  91 effectiveStdenv.mkDerivation rec {
  92   pname = "onnxruntime";
  93   inherit version;
  94
  95   src = fetchFromGitHub {
  96     owner = "microsoft";
  97     repo = "onnxruntime";
  98     rev = "refs/tags/v${version}";
  99     hash = "sha256-+zWtbLKekGhwdBU3bm1u2F7rYejQ62epE+HcHj05/8A=";
 100     fetchSubmodules = true;
 101   };
 102
 103   patches = [
 104     # If you stumble on these patches trying to update onnxruntime, check
 105     # `git blame` and ping the introducers.
 106
 107     # Context: we want the upstream to
 108     # - always try find_package first (FIND_PACKAGE_ARGS),
 109     # - use MakeAvailable instead of the low-level Populate,
 110     # - use Eigen3::Eigen as the target name (as declared by libeigen/eigen).
 111     ./0001-eigen-allow-dependency-injection.patch
 112     # Incorporate a patch that has landed upstream which exposes new
 113     # 'abseil-cpp' libraries & modifies the 're2' CMakeLists to fix a
 114     # configuration error that around missing 'gmock' exports.
 115     #
 116     # TODO: Check if it can be dropped after 1.19.0
 117     # https://github.com/microsoft/onnxruntime/commit/b522df0ae477e59f60acbe6c92c8a64eda96cace
 118     ./update-re2.patch
 119   ] ++ lib.optionals cudaSupport [
 120     # We apply the referenced 1064.patch ourselves to our nix dependency.
 121     #  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
 122     ./nvcc-gsl.patch
 123   ];
 124
 125   nativeBuildInputs = [
 126     cmake
 127     pkg-config
 128     python3Packages.python
 129     protobuf_21
 130   ] ++ lib.optionals pythonSupport (with python3Packages; [
 131     pip
 132     python
 133     pythonOutputDistHook
 134     setuptools
 135     wheel
 136   ]) ++ lib.optionals cudaSupport [
 137     cudaPackages.cuda_nvcc
 138   ];
 139
 140   buildInputs = [
 141     cpuinfo
 142     eigen
 143     glibcLocales
 144     libpng
 145     nlohmann_json
 146     microsoft-gsl
 147     pytorch_clog
 148     zlib
 149   ] ++ lib.optionals pythonSupport (with python3Packages; [
 150     numpy
 151     pybind11
 152     packaging
 153   ]) ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [
 154     Foundation
 155     libiconv
 156   ] ++ lib.optionals cudaSupport (with cudaPackages; [
 157     cuda_cccl # cub/cub.cuh
 158     libcublas # cublas_v2.h
 159     libcurand # curand.h
 160     libcusparse # cusparse.h
 161     libcufft # cufft.h
 162     cudnn # cudnn.h
 163     cuda_cudart
 164   ] ++ lib.optionals (cudaSupport && ncclSupport) (with cudaPackages; [
 165     nccl
 166   ]));
 167
 168   nativeCheckInputs = [
 169     gtest
 170   ] ++ lib.optionals pythonSupport (with python3Packages; [
 171     pytest
 172     sympy
 173     onnx
 174   ]);
 175
 176   # TODO: build server, and move .so's to lib output
 177   # Python's wheel is stored in a separate dist output
 178   outputs = [ "out" "dev" ] ++ lib.optionals pythonSupport [ "dist" ];
 179
 180   enableParallelBuilding = true;
 181
 182   cmakeDir = "../cmake";
 183
 184   cmakeFlags = [
 185     "-DABSL_ENABLE_INSTALL=ON"
 186     "-DFETCHCONTENT_FULLY_DISCONNECTED=ON"
 187     "-DFETCHCONTENT_QUIET=OFF"
 188     "-DFETCHCONTENT_SOURCE_DIR_ABSEIL_CPP=${abseil-cpp.src}"
 189     "-DFETCHCONTENT_SOURCE_DIR_DATE=${howard-hinnant-date}"
 190     "-DFETCHCONTENT_SOURCE_DIR_FLATBUFFERS=${flatbuffers_23.src}"
 191     "-DFETCHCONTENT_SOURCE_DIR_GOOGLETEST=${gtest.src}"
 192     "-DFETCHCONTENT_SOURCE_DIR_GOOGLE_NSYNC=${nsync.src}"
 193     "-DFETCHCONTENT_SOURCE_DIR_MP11=${mp11}"
 194     "-DFETCHCONTENT_SOURCE_DIR_ONNX=${onnx}"
 195     "-DFETCHCONTENT_SOURCE_DIR_RE2=${re2.src}"
 196     "-DFETCHCONTENT_SOURCE_DIR_SAFEINT=${safeint}"
 197     "-DFETCHCONTENT_TRY_FIND_PACKAGE_MODE=ALWAYS"
 198     "-Donnxruntime_BUILD_SHARED_LIB=ON"
 199     (lib.cmakeBool "onnxruntime_BUILD_UNIT_TESTS" doCheck)
 200     "-Donnxruntime_ENABLE_LTO=ON"
 201     "-Donnxruntime_USE_FULL_PROTOBUF=OFF"
 202     (lib.cmakeBool "onnxruntime_USE_CUDA" cudaSupport)
 203     (lib.cmakeBool "onnxruntime_USE_NCCL" (cudaSupport && ncclSupport))
 204   ] ++ lib.optionals pythonSupport [
 205     "-Donnxruntime_ENABLE_PYTHON=ON"
 206   ] ++ lib.optionals cudaSupport [
 207     (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${cutlass}")
 208     (lib.cmakeFeature "onnxruntime_CUDNN_HOME" "${cudaPackages.cudnn}")
 209     (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaArchitecturesString)
 210     (lib.cmakeFeature "onnxruntime_NVCC_THREADS" "1")
 211   ];
 212
 213   env = lib.optionalAttrs effectiveStdenv.cc.isClang {
 214     NIX_CFLAGS_COMPILE = toString [
 215       "-Wno-error=deprecated-declarations"
 216       "-Wno-error=deprecated-pragma"
 217       "-Wno-error=unused-but-set-variable"
 218     ];
 219   };
 220
 221   # aarch64-linux fails cpuinfo test, because /sys/devices/system/cpu/ does not exist in the sandbox
 222   doCheck = !(cudaSupport || effectiveStdenv.buildPlatform.system == "aarch64-linux");
 223
 224   requiredSystemFeatures = lib.optionals cudaSupport [ "big-parallel" ];
 225
 226   postPatch = ''
 227     substituteInPlace cmake/libonnxruntime.pc.cmake.in \
 228       --replace-fail '$'{prefix}/@CMAKE_INSTALL_ @CMAKE_INSTALL_
 229   '' + lib.optionalString (effectiveStdenv.hostPlatform.system == "aarch64-linux") ''
 230     # https://github.com/NixOS/nixpkgs/pull/226734#issuecomment-1663028691
 231     rm -v onnxruntime/test/optimizer/nhwc_transformer_test.cc
 232   '';
 233
 234   postBuild = lib.optionalString pythonSupport ''
 235     ${python3Packages.python.interpreter} ../setup.py bdist_wheel
 236   '';
 237
 238   postInstall = ''
 239     # perform parts of `tools/ci_build/github/linux/copy_strip_binary.sh`
 240     install -m644 -Dt $out/include \
 241       ../include/onnxruntime/core/framework/provider_options.h \
 242       ../include/onnxruntime/core/providers/cpu/cpu_provider_factory.h \
 243       ../include/onnxruntime/core/session/onnxruntime_*.h
 244   '';
 245
 246   passthru = {
 247     inherit cudaSupport cudaPackages; # for the python module
 248     protobuf = protobuf_21;
 249     tests = lib.optionalAttrs pythonSupport {
 250       python = python3Packages.onnxruntime;
 251     };
 252   };
 253
 254   meta = with lib; {
 255     description = "Cross-platform, high performance scoring engine for ML models";
 256     longDescription = ''
 257       ONNX Runtime is a performance-focused complete scoring engine
 258       for Open Neural Network Exchange (ONNX) models, with an open
 259       extensible architecture to continually address the latest developments
 260       in AI and Deep Learning. ONNX Runtime stays up to date with the ONNX
 261       standard with complete implementation of all ONNX operators, and
 262       supports all ONNX releases (1.2+) with both future and backwards
 263       compatibility.
 264     '';
 265     homepage = "https://github.com/microsoft/onnxruntime";
 266     changelog = "https://github.com/microsoft/onnxruntime/releases/tag/v${version}";
 267     # https://github.com/microsoft/onnxruntime/blob/master/BUILD.md#architectures
 268     platforms = platforms.unix;
 269     license = licenses.mit;
 270     maintainers = with maintainers; [ puffnfresh ck3d cbourjau ];
 271   };
 272 }