pkgs/development/libraries/onnxruntime/default.nix

   1 { stdenv
   2 , lib
   3 , fetchFromGitHub
   4 , fetchpatch
   5 , fetchurl
   6 , pkg-config
   7 , cmake
   8 , python3Packages
   9 , libpng
  10 , zlib
  11 , eigen
  12 , protobuf
  13 , howard-hinnant-date
  14 , nlohmann_json
  15 , boost
  16 , oneDNN
  17 , gtest
  18 , pythonSupport ? false
  19 , nsync
  20 , flatbuffers
  21 }:
  22
  23 # Python Support
  24 #
  25 # When enabling Python support a wheel is made and stored in a `dist` output.
  26 # This wheel is then installed in a separate derivation.
  27
  28 assert pythonSupport -> lib.versionOlder protobuf.version "3.20";
  29
  30 let
  31   # prefetch abseil
  32   # Note: keep URL in sync with `cmake/external/abseil-cpp.cmake`
  33   abseil = fetchurl {
  34     url = "https://github.com/abseil/abseil-cpp/archive/refs/tags/20211102.0.zip";
  35     sha256 = "sha256-pFZ/8C+spnG5XjHTFbqxi0K2xvGmDpHG6oTlohQhEsI=";
  36   };
  37 in
  38 stdenv.mkDerivation rec {
  39   pname = "onnxruntime";
  40   version = "1.12.1";
  41
  42   src = fetchFromGitHub {
  43     owner = "microsoft";
  44     repo = "onnxruntime";
  45     rev = "v${version}";
  46     sha256 = "sha256-wwllEemiHTp9aJcCd1gsTS4WUVMp5wW+4i/+6DzmAeM=";
  47     fetchSubmodules = true;
  48   };
  49
  50   patches = [
  51     # Use dnnl from nixpkgs instead of submodules
  52     (fetchpatch {
  53       name = "system-dnnl.patch";
  54       url = "https://aur.archlinux.org/cgit/aur.git/plain/system-dnnl.diff?h=python-onnxruntime&id=0185531906bda3a9aba93bbb0f3dcfeb0ae671ad";
  55       sha256 = "sha256-58RBrQnAWNtc/1pmFs+PkZ6qCsL1LfMY3P0exMKzotA=";
  56     })
  57   ];
  58
  59   nativeBuildInputs = [
  60     cmake
  61     pkg-config
  62     python3Packages.python
  63     gtest
  64   ] ++ lib.optionals pythonSupport (with python3Packages; [
  65     setuptools
  66     wheel
  67     pip
  68     pythonOutputDistHook
  69   ]);
  70
  71   buildInputs = [
  72     libpng
  73     zlib
  74     howard-hinnant-date
  75     nlohmann_json
  76     boost
  77     oneDNN
  78     protobuf
  79   ] ++ lib.optionals pythonSupport [
  80     nsync
  81     python3Packages.numpy
  82     python3Packages.pybind11
  83   ];
  84
  85   # TODO: build server, and move .so's to lib output
  86   # Python's wheel is stored in a separate dist output
  87   outputs = [ "out" "dev" ] ++ lib.optionals pythonSupport [ "dist" ];
  88
  89   enableParallelBuilding = true;
  90
  91   cmakeDir = "../cmake";
  92
  93   cmakeFlags = [
  94     "-Donnxruntime_PREFER_SYSTEM_LIB=ON"
  95     "-Donnxruntime_BUILD_SHARED_LIB=ON"
  96     "-Donnxruntime_ENABLE_LTO=ON"
  97     "-Donnxruntime_BUILD_UNIT_TESTS=ON"
  98     "-Donnxruntime_USE_PREINSTALLED_EIGEN=ON"
  99     "-Donnxruntime_USE_MPI=ON"
 100     "-Deigen_SOURCE_PATH=${eigen.src}"
 101     "-Donnxruntime_USE_DNNL=YES"
 102   ] ++ lib.optionals pythonSupport [
 103     "-Donnxruntime_ENABLE_PYTHON=ON"
 104   ];
 105
 106   doCheck = true;
 107
 108   postPatch = ''
 109     substituteInPlace cmake/external/abseil-cpp.cmake \
 110       --replace "${abseil.url}" "${abseil}"
 111
 112     substituteInPlace cmake/libonnxruntime.pc.cmake.in \
 113       --replace '$'{prefix}/@CMAKE_INSTALL_ @CMAKE_INSTALL_
 114   '';
 115
 116   postBuild = lib.optionalString pythonSupport ''
 117     ${python3Packages.python.interpreter} ../setup.py bdist_wheel
 118   '';
 119
 120   postInstall = ''
 121     # perform parts of `tools/ci_build/github/linux/copy_strip_binary.sh`
 122     install -m644 -Dt $out/include \
 123       ../include/onnxruntime/core/framework/provider_options.h \
 124       ../include/onnxruntime/core/providers/cpu/cpu_provider_factory.h \
 125       ../include/onnxruntime/core/session/onnxruntime_*.h
 126   '';
 127
 128   passthru = {
 129     inherit protobuf;
 130     tests = lib.optionalAttrs pythonSupport {
 131       python = python3Packages.onnxruntime;
 132     };
 133   };
 134
 135   meta = with lib; {
 136     description = "Cross-platform, high performance scoring engine for ML models";
 137     longDescription = ''
 138       ONNX Runtime is a performance-focused complete scoring engine
 139       for Open Neural Network Exchange (ONNX) models, with an open
 140       extensible architecture to continually address the latest developments
 141       in AI and Deep Learning. ONNX Runtime stays up to date with the ONNX
 142       standard with complete implementation of all ONNX operators, and
 143       supports all ONNX releases (1.2+) with both future and backwards
 144       compatibility.
 145     '';
 146     homepage = "https://github.com/microsoft/onnxruntime";
 147     changelog = "https://github.com/microsoft/onnxruntime/releases/tag/v${version}";
 148     # https://github.com/microsoft/onnxruntime/blob/master/BUILD.md#architectures
 149     platforms = platforms.unix;
 150     license = licenses.mit;
 151     maintainers = with maintainers; [ jonringer puffnfresh ck3d ];
 152   };
 153 }