onnx-simplifier

1
from distutils.spawn import find_executable
2
from distutils import sysconfig, log
3
import setuptools
4
import setuptools.command.build_py
5
import setuptools.command.develop
6
import setuptools.command.build_ext
7

8
from collections import namedtuple
9
from contextlib import contextmanager
10
import glob
11
import os
12
import shlex
13
import subprocess
14
import sys
15
import platform
16
from textwrap import dedent
17
import multiprocessing
18
import re
19

20

21
TOP_DIR = os.path.realpath(os.path.dirname(__file__))
22
SRC_DIR = os.path.join(TOP_DIR, 'onnxsim')
23
CMAKE_BUILD_DIR = os.path.join(TOP_DIR, '.setuptools-cmake-build')
24

25
WINDOWS = (os.name == 'nt')
26
MACOS = sys.platform.startswith("darwin")
27

28
CMAKE = find_executable('cmake')
29

30
install_requires = []
31
setup_requires = []
32

33
USE_MSVC_STATIC_RUNTIME = bool(os.getenv('USE_MSVC_STATIC_RUNTIME', '0') == '1')
34
ONNX_ML = not bool(os.getenv('ONNX_ML') == '0')
35
ONNX_VERIFY_PROTO3 = bool(os.getenv('ONNX_VERIFY_PROTO3') == '1')
36
ONNX_NAMESPACE = os.getenv('ONNX_NAMESPACE', 'onnx')
37
ONNX_BUILD_TESTS = bool(os.getenv('ONNX_BUILD_TESTS') == '1')
38
ONNX_OPT_USE_SYSTEM_PROTOBUF = bool(os.getenv('ONNX_OPT_USE_SYSTEM_PROTOBUF', '0') == '1')
39

40
DEBUG = bool(os.getenv('DEBUG'))
41
COVERAGE = bool(os.getenv('COVERAGE'))
42

43
try:
44
    version = subprocess.check_output(['git', 'describe', '--tags', '--abbrev=0'],
45
                                      cwd=TOP_DIR).decode('ascii').strip()
46
    if version[0] == 'v':
47
        version = version[1:]
48
except (OSError, subprocess.CalledProcessError):
49
    with open(os.path.join(TOP_DIR, 'VERSION')) as ver_file:
50
        version = ver_file.read().strip()
51

52
try:
53
    git_version = subprocess.check_output(['git', 'rev-parse', 'HEAD'],
54
                                          cwd=TOP_DIR).decode('ascii').strip()
55
except (OSError, subprocess.CalledProcessError):
56
    git_version = None
57

58
if os.getenv('ONNXSIM_SDIST') is not None:
59
    version = '0.0.0'
60
    git_version = None
61

62
VersionInfo = namedtuple('VersionInfo', ['version', 'git_version'])(
63
    version=version,
64
    git_version=git_version
65
)
66

67
assert CMAKE, 'Could not find "cmake" executable!'
68

69
@contextmanager
70
def cd(path):
71
    if not os.path.isabs(path):
72
        raise RuntimeError('Can only cd to absolute path, got: {}'.format(path))
73
    orig_path = os.getcwd()
74
    os.chdir(path)
75
    try:
76
        yield
77
    finally:
78
        os.chdir(orig_path)
79

80

81
class ONNXCommand(setuptools.Command):
82
    user_options = []
83

84
    def initialize_options(self):
85
        pass
86

87
    def finalize_options(self):
88
        pass
89

90

91
class create_version(ONNXCommand):
92
    def run(self):
93
        with open(os.path.join(SRC_DIR, 'version.py'), 'w') as f:
94
            f.write(dedent('''\
95
            # This file is generated by setup.py. DO NOT EDIT!
96

97
            version = '{version}'
98
            git_version = '{git_version}'
99
            '''.format(**dict(VersionInfo._asdict()))))
100

101

102
class cmake_build(setuptools.Command):
103
    """
104
    Compiles everything when `python setupmnm.py build` is run using cmake.
105

106
    Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
107
    environment variable.
108

109
    The number of CPUs used by `make` can be specified by passing `-j<ncpus>`
110
    to `setup.py build`.  By default all CPUs are used.
111
    """
112
    user_options = [
113
        (str('jobs='), str('j'), str('Specifies the number of jobs to use with make'))
114
    ]
115

116
    built = False
117

118
    def initialize_options(self):
119
        self.jobs = None
120

121
    def finalize_options(self):
122
        self.set_undefined_options('build', ('parallel', 'jobs'))
123
        if self.jobs is None and os.getenv("MAX_JOBS") is not None:
124
            self.jobs = os.getenv("MAX_JOBS")
125
        self.jobs = multiprocessing.cpu_count() if self.jobs is None else int(self.jobs)
126

127
    def run(self):
128
        if cmake_build.built:
129
            return
130
        cmake_build.built = True
131
        if not os.path.exists(CMAKE_BUILD_DIR):
132
            os.makedirs(CMAKE_BUILD_DIR)
133

134
        with cd(CMAKE_BUILD_DIR):
135
            build_type = 'Release'
136
            # configure
137
            cmake_args = [
138
                CMAKE,
139
                '-DPython_INCLUDE_DIR={}'.format(sysconfig.get_python_inc()),
140
                '-DPython_EXECUTABLE={}'.format(sys.executable),
141
                # For pybind11
142
                '-DPYTHON_EXECUTABLE={}'.format(sys.executable),
143
                '-DBUILD_ONNX_PYTHON=OFF',
144
                '-DONNXSIM_PYTHON=ON',
145
                '-DONNXSIM_BUILTIN_ORT=OFF',
146
                '-DONNX_USE_LITE_PROTO=OFF',
147
                '-DCMAKE_EXPORT_COMPILE_COMMANDS=ON',
148
                '-DONNX_NAMESPACE={}'.format(ONNX_NAMESPACE),
149
                '-DPY_EXT_SUFFIX={}'.format(
150
                    sysconfig.get_config_var('EXT_SUFFIX') or ''),
151
                '-DONNX_OPT_USE_SYSTEM_PROTOBUF={}'.format(
152
                    'ON' if ONNX_OPT_USE_SYSTEM_PROTOBUF else 'OFF'),
153
            ]
154
            if COVERAGE:
155
                cmake_args.append('-DONNX_COVERAGE=ON')
156
            if COVERAGE or DEBUG:
157
                # in order to get accurate coverage information, the
158
                # build needs to turn off optimizations
159
                build_type = 'Debug'
160
            cmake_args.append('-DCMAKE_BUILD_TYPE=%s' % build_type)
161
            if WINDOWS:
162
                cmake_args.extend([
163
                    # we need to link with libpython on windows, so
164
                    # passing python version to window in order to
165
                    # find python in cmake
166
                    '-DPY_VERSION={}'.format('{0}.{1}'.format(* \
167
                                                              sys.version_info[:2])),
168
                ])
169
                if USE_MSVC_STATIC_RUNTIME:
170
                    cmake_args.append('-DONNX_USE_MSVC_STATIC_RUNTIME=ON')
171
                if platform.architecture()[0] == '64bit':
172
                    cmake_args.extend(['-A', 'x64', '-T', 'host=x64'])
173
                else:
174
                    cmake_args.extend(['-A', 'Win32', '-T', 'host=x86'])
175
            if MACOS:
176
                # Cross-compile support for macOS - respect ARCHFLAGS if set
177
                archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
178
                if archs:
179
                    cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
180
            if ONNX_ML:
181
                cmake_args.append('-DONNX_ML=1')
182
            if ONNX_VERIFY_PROTO3:
183
                cmake_args.append('-DONNX_VERIFY_PROTO3=1')
184
            if ONNX_BUILD_TESTS:
185
                cmake_args.append('-DONNX_BUILD_TESTS=ON')
186
            if 'CMAKE_ARGS' in os.environ:
187
                extra_cmake_args = shlex.split(os.environ['CMAKE_ARGS'])
188
                # prevent crossfire with downstream scripts
189
                del os.environ['CMAKE_ARGS']
190
                log.info('Extra cmake args: {}'.format(extra_cmake_args))
191
                cmake_args.extend(extra_cmake_args)
192
            cmake_args.append(TOP_DIR)
193
            print(f"Run command {cmake_args}")
194
            subprocess.check_call(cmake_args)
195

196
            build_args = [CMAKE, '--build', os.curdir, '--target onnxsim_cpp2py_export']
197
            if WINDOWS:
198
                build_args.extend(['--config', build_type])
199
                build_args.extend(['--', '/maxcpucount:{}'.format(self.jobs)])
200
            else:
201
                build_args.extend(['--', '-j', str(self.jobs)])
202
            print(f"Run command {build_args}")
203
            subprocess.check_call(build_args)
204

205

206
class build_py(setuptools.command.build_py.build_py):
207
    def run(self):
208
        self.run_command('create_version')
209
        return setuptools.command.build_py.build_py.run(self)
210

211

212
class develop(setuptools.command.develop.develop):
213
    def run(self):
214
        self.run_command('build_py')
215
        setuptools.command.develop.develop.run(self)
216

217

218
class build_ext(setuptools.command.build_ext.build_ext):
219
    def run(self):
220
        self.run_command('cmake_build')
221
        setuptools.command.build_ext.build_ext.run(self)
222

223
    def build_extensions(self):
224
        for ext in self.extensions:
225
            fullname = self.get_ext_fullname(ext.name)
226
            filename = os.path.basename(self.get_ext_filename(fullname))
227

228
            lib_path = CMAKE_BUILD_DIR
229
            if os.name == 'nt':
230
                debug_lib_dir = os.path.join(lib_path, "Debug")
231
                release_lib_dir = os.path.join(lib_path, "Release")
232
                if os.path.exists(debug_lib_dir):
233
                    lib_path = debug_lib_dir
234
                elif os.path.exists(release_lib_dir):
235
                    lib_path = release_lib_dir
236
            src = os.path.join(lib_path, filename)
237
            dst_dir = os.path.join(os.path.realpath(
238
                self.build_lib), "onnxsim")
239
            dst = os.path.join(dst_dir, filename)
240
            os.makedirs(dst_dir, exist_ok=True)
241
            self.copy_file(src, dst)
242

243

244
cmdclass = {
245
    'create_version': create_version,
246
    'cmake_build': cmake_build,
247
    'build_ext': build_ext,
248
    'build_py': build_py,
249
    'develop': develop,
250
}
251

252
ext_modules = [
253
    setuptools.Extension(
254
        name=str('onnxsim.onnxsim_cpp2py_export'),
255
        sources=[])
256
]
257

258
# no need to do fancy stuff so far
259
packages = setuptools.find_packages()
260

261
# Though we depend on onnxruntime, it has three different packages:
262
# onnxruntime, onnxruntime-gpu and onnxruntime-noopenmp.
263
# The solution is, we publish two packages, a wheel named onnxsim-no-ort
264
# and a sdist package named onnxsim, onnxsim depends on onnxsim-no-ort,
265
# and also check if one of onnxruntime packages is installed, and depends
266
# on onnxruntime when no existing installed packages.
267
install_requires.extend([
268
    'onnx',
269
    'rich',
270
])
271

272
setup_requires.append('pytest-runner')
273

274
# read the contents of your README file
275
from pathlib import Path
276
this_directory = Path(__file__).parent
277
long_description = (this_directory / "README.md").read_text()
278

279
setuptools.setup(
280
    name=os.getenv("ONNXSIM_PKG_NAME", "onnxsim"),
281
    version=VersionInfo.version,
282
    description='Simplify your ONNX model',
283
    ext_modules=ext_modules,
284
    cmdclass=cmdclass,
285
    packages=packages,
286
    license='Apache License v2.0',
287
    include_package_data=True,
288
    install_requires=install_requires,
289
    setup_requires=setup_requires,
290
    author='ONNX Simplifier Authors',
291
    author_email='daquexian566@gmail.com',
292
    url='https://github.com/daquexian/onnx-simplifier',
293
    keywords='deep-learning ONNX',
294
    long_description=long_description,
295
    long_description_content_type='text/markdown',
296
    classifiers=[
297
        'Development Status :: 4 - Beta',
298
        'Intended Audience :: Developers',
299
        'License :: OSI Approved :: Apache Software License',
300
        'Programming Language :: Python :: 3 :: Only',
301
        'Programming Language :: Python :: 3.7',
302
        'Programming Language :: Python :: 3.8',
303
        'Programming Language :: Python :: 3.9',
304
        'Programming Language :: Python :: 3.10',
305
        'Programming Language :: Python :: 3.11',
306
        'Topic :: Scientific/Engineering',
307
        'Topic :: Software Development'
308
    ],
309
    python_requires='>=3.7',
310
    entry_points={
311
        'console_scripts': [
312
            'onnxsim=onnxsim:main',
313
        ],
314
    },
315
)
316