summaryrefslogtreecommitdiff
path: root/var/spack/repos/builtin/packages/py-onnx-runtime/package.py
blob: 2bb30786a84060c0cab6d750445acde1b671124e (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *


class PyOnnxRuntime(CMakePackage, PythonPackage):
    """ONNX Runtime is a performance-focused complete scoring
    engine for Open Neural Network Exchange (ONNX) models, with
    an open extensible architecture to continually address the
    latest developments in AI and Deep Learning. ONNX Runtime
    stays up to date with the ONNX standard with complete
    implementation of all ONNX operators, and supports all
    ONNX releases (1.2+) with both future and backwards
    compatibility."""

    homepage = "https://github.com/microsoft/onnxruntime"
    git      = "https://github.com/microsoft/onnxruntime.git"

    version('1.7.2',  tag='v1.7.2',  submodules=True)

    variant('cuda', default=False, description='Build with CUDA support')

    depends_on('cmake@3.1:', type='build')
    depends_on('ninja', type='build')
    depends_on('python', type=('build', 'run'))
    depends_on('protobuf')
    depends_on('py-protobuf', type=('build', 'run'))
    depends_on('py-setuptools', type='build')
    depends_on('py-numpy@1.16.6:', type=('build', 'run'))
    depends_on('py-wheel', type='build')
    depends_on('py-onnx', type=('build', 'run'))
    depends_on('zlib')
    depends_on('libpng')
    depends_on('py-pybind11', type='build')
    depends_on('cuda', when='+cuda')
    depends_on('cudnn', when='+cuda')
    depends_on('iconv', type=('build', 'link', 'run'))
    depends_on('re2+shared')

    extends('python')
    # Adopted from CMS experiment's fork of onnxruntime
    # https://github.com/cms-externals/onnxruntime/compare/5bc92df...d594f80
    patch('cms.patch', level=1, when='@1.7.2')
    # https://github.com/microsoft/onnxruntime/issues/4234#issuecomment-698077636
    patch('libiconv.patch', level=0, when='@1.7.2')
    # https://github.com/microsoft/onnxruntime/commit/de4089f8cbe0baffe56a363cc3a41595cc8f0809.patch
    patch('gcc11.patch', level=1, when='@1.7.2')

    dynamic_cpu_arch_values = ('NOAVX', 'AVX', 'AVX2', 'AVX512')

    variant('dynamic_cpu_arch', default='AVX512',
            values=dynamic_cpu_arch_values, multi=False,
            description='AVX support level')

    generator = 'Ninja'
    root_cmakelists_dir = 'cmake'

    def setup_build_environment(self, env):
        value = self.spec.variants['dynamic_cpu_arch'].value
        value = self.dynamic_cpu_arch_values.index(value)
        env.set('MLAS_DYNAMIC_CPU_ARCH', str(value))

    def setup_run_environment(self, env):
        value = self.spec.variants['dynamic_cpu_arch'].value
        value = self.dynamic_cpu_arch_values.index(value)
        env.set('MLAS_DYNAMIC_CPU_ARCH', str(value))

    def cmake_args(self):
        define = self.define
        define_from_variant = self.define_from_variant

        args = [define('onnxruntime_ENABLE_PYTHON', True),
                define('onnxruntime_BUILD_SHARED_LIB', True),
                define_from_variant('onnxruntime_USE_CUDA', 'cuda'),
                define('onnxruntime_BUILD_CSHARP', False),
                define('onnxruntime_USE_EIGEN_FOR_BLAS', True),
                define('onnxruntime_USE_OPENBLAS', False),
                define("onnxruntime_USE_MKLML", False),
                define("onnxruntime_USE_NGRAPH", False),
                define("onnxruntime_USE_OPENMP", False),
                define("onnxruntime_USE_TVM", False),
                define("onnxruntime_USE_LLVM", False),
                define("onnxruntime_ENABLE_MICROSOFT_INTERNAL", False),
                define("onnxruntime_USE_BRAINSLICE", False),
                define("onnxruntime_USE_NUPHAR", False),
                define("onnxruntime_USE_TENSORRT", False),
                define("onnxruntime_CROSS_COMPILING", False),
                define("onnxruntime_USE_FULL_PROTOBUF", True),
                define("onnxruntime_DISABLE_CONTRIB_OPS", False),
                define("onnxruntime_USE_PREINSTALLED_PROTOBUF", True),
                define("onnxruntime_PREFER_SYSTEM_LIB", True)]

        if self.spec.satisfies('+cuda'):
            args.extend((
                define('onnxruntime_CUDA_VERSION', str(self.spec['cuda'].version)),
                define('onnxruntime_CUDA_HOME', self.spec['cuda'].prefix),
                define('onnxruntime_CUDNN_HOME', self.spec['cudnn'].prefix),
                define('CMAKE_CUDA_FLAGS', '-cudart shared'),
                define('CMAKE_CUDA_RUNTIME_LIBRARY', 'Shared'),
                define('DCMAKE_TRY_COMPILE_PLATFORM_VARIABLES',
                       'CMAKE_CUDA_RUNTIME_LIBRARY')
            ))

        return args

    def setup_file(self):
        return join_path(self.stage.source_path, 'setup.py')

    @run_after('build')
    def build_python(self):
        """Build everything needed to install."""
        with working_dir(self.stage.source_path):
            PythonPackage.build(self, self.spec, self.prefix)

    @run_after('install')
    def install_python(self):
        with working_dir(self.stage.source_path):
            PythonPackage.install(self, self.spec, self.prefix)