Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 2 | # Copyright © 2020 Arm Ltd. All rights reserved. |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 3 | # Copyright © 2020 NXP and Contributors. All rights reserved. |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 4 | # SPDX-License-Identifier: MIT |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 5 | """Python bindings for Arm NN |
| 6 | |
| 7 | PyArmNN is a python extension for Arm NN SDK providing an interface similar to Arm NN C++ API. |
| 8 | """ |
| 9 | __version__ = None |
| 10 | __arm_ml_version__ = None |
| 11 | |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 12 | import logging |
| 13 | import os |
| 14 | import sys |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 15 | import subprocess |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 16 | from functools import lru_cache |
| 17 | from pathlib import Path |
| 18 | from itertools import chain |
| 19 | |
| 20 | from setuptools import setup |
| 21 | from distutils.core import Extension |
| 22 | from setuptools.command.build_py import build_py |
| 23 | from setuptools.command.build_ext import build_ext |
| 24 | |
| 25 | logger = logging.Logger(__name__) |
| 26 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 27 | DOCLINES = __doc__.split("\n") |
| 28 | LIB_ENV_NAME = "ARMNN_LIB" |
| 29 | INCLUDE_ENV_NAME = "ARMNN_INCLUDE" |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 30 | |
| 31 | |
| 32 | def check_armnn_version(*args): |
| 33 | pass |
| 34 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 35 | __current_dir = os.path.dirname(os.path.realpath(__file__)) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 36 | |
Kevin May | 8acafd1 | 2020-06-02 11:06:46 +0100 | [diff] [blame] | 37 | exec(open(os.path.join(__current_dir, 'src', 'pyarmnn', '_version.py'), encoding="utf-8").read()) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 38 | |
| 39 | |
| 40 | class ExtensionPriorityBuilder(build_py): |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 41 | """Runs extension builder before other stages. Otherwise generated files are not included to the distribution. |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 42 | """ |
| 43 | |
| 44 | def run(self): |
| 45 | self.run_command('build_ext') |
| 46 | return super().run() |
| 47 | |
| 48 | |
| 49 | class ArmnnVersionCheckerExtBuilder(build_ext): |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 50 | """Builds an extension (i.e. wrapper). Additionally checks for version. |
| 51 | """ |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 52 | |
| 53 | def __init__(self, dist): |
| 54 | super().__init__(dist) |
| 55 | self.failed_ext = [] |
| 56 | |
| 57 | def build_extension(self, ext): |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 58 | if ext.optional: |
| 59 | try: |
| 60 | super().build_extension(ext) |
| 61 | except Exception as err: |
| 62 | self.failed_ext.append(ext) |
| 63 | logger.warning('Failed to build extension %s. \n %s', ext.name, str(err)) |
| 64 | else: |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 65 | super().build_extension(ext) |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 66 | if ext.name == 'pyarmnn._generated._pyarmnn_version': |
| 67 | sys.path.append(os.path.abspath(os.path.join(self.build_lib, str(Path(ext._file_name).parent)))) |
| 68 | from _pyarmnn_version import GetVersion |
| 69 | check_armnn_version(GetVersion(), __arm_ml_version__) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 70 | |
| 71 | def copy_extensions_to_source(self): |
| 72 | |
| 73 | for ext in self.failed_ext: |
| 74 | self.extensions.remove(ext) |
| 75 | super().copy_extensions_to_source() |
| 76 | |
| 77 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 78 | def linux_gcc_name(): |
| 79 | """Returns the name of the `gcc` compiler. Might happen that we are cross-compiling and the |
| 80 | compiler has a longer name. |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 81 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 82 | Args: |
| 83 | None |
| 84 | |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 85 | Returns: |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 86 | str: Name of the `gcc` compiler or None |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 87 | """ |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 88 | cc_env = os.getenv('CC') |
| 89 | if cc_env is not None: |
| 90 | if subprocess.Popen([cc_env, "--version"], stdout=subprocess.DEVNULL): |
| 91 | return cc_env |
| 92 | return "gcc" if subprocess.Popen(["gcc", "--version"], stdout=subprocess.DEVNULL) else None |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 93 | |
| 94 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 95 | def linux_gcc_lib_search(gcc_compiler_name: str = linux_gcc_name()): |
| 96 | """Calls the `gcc` to get linker default system paths. |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 97 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 98 | Args: |
| 99 | gcc_compiler_name(str): Name of the GCC compiler |
| 100 | |
| 101 | Returns: |
| 102 | list: A list of paths. |
| 103 | |
| 104 | Raises: |
| 105 | RuntimeError: If unable to find GCC. |
| 106 | """ |
| 107 | if gcc_compiler_name is None: |
| 108 | raise RuntimeError("Unable to find gcc compiler") |
| 109 | cmd1 = subprocess.Popen([gcc_compiler_name, "--print-search-dirs"], stdout=subprocess.PIPE) |
| 110 | cmd2 = subprocess.Popen(["grep", "libraries"], stdin=cmd1.stdout, |
| 111 | stdout=subprocess.PIPE, stderr=subprocess.DEVNULL) |
| 112 | cmd1.stdout.close() |
| 113 | out, _ = cmd2.communicate() |
| 114 | out = out.decode("utf-8").split('=') |
| 115 | return tuple(out[1].split(':')) if len(out) > 0 else None |
| 116 | |
| 117 | |
| 118 | def find_includes(armnn_include_env: str = INCLUDE_ENV_NAME): |
| 119 | """Searches for ArmNN includes. |
| 120 | |
| 121 | Args: |
| 122 | armnn_include_env(str): Environmental variable to use as path. |
| 123 | |
| 124 | Returns: |
| 125 | list: A list of paths to include. |
| 126 | """ |
| 127 | armnn_include_path = os.getenv(armnn_include_env) |
| 128 | if armnn_include_path is not None and os.path.exists(armnn_include_path): |
| 129 | armnn_include_path = [armnn_include_path] |
| 130 | else: |
| 131 | armnn_include_path = ['/usr/local/include', '/usr/include'] |
| 132 | return armnn_include_path |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 133 | |
| 134 | |
| 135 | @lru_cache(maxsize=1) |
| 136 | def find_armnn(lib_name: str, |
| 137 | optional: bool = False, |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 138 | armnn_libs_env: str = LIB_ENV_NAME, |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 139 | default_lib_search: tuple = linux_gcc_lib_search()): |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 140 | """Searches for ArmNN installation on the local machine. |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 141 | |
| 142 | Args: |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 143 | lib_name(str): Lib name to find. |
| 144 | optional(bool): Do not fail if optional. Default is False - fail if library was not found. |
| 145 | armnn_libs_env(str): Custom environment variable pointing to ArmNN libraries location, default is 'ARMNN_LIBS' |
| 146 | default_lib_search(tuple): list of paths to search for ArmNN if not found within path provided by 'ARMNN_LIBS' |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 147 | env variable |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 148 | Returns: |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 149 | tuple: Contains name of the armnn libs, paths to the libs. |
| 150 | |
| 151 | Raises: |
| 152 | RuntimeError: If armnn libs are not found. |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 153 | """ |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 154 | armnn_lib_path = os.getenv(armnn_libs_env) |
| 155 | lib_search = [armnn_lib_path] if armnn_lib_path is not None else default_lib_search |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 156 | armnn_libs = dict(map(lambda path: (':{}'.format(path.name), path), |
| 157 | chain.from_iterable(map(lambda lib_path: Path(lib_path).glob(lib_name), |
| 158 | lib_search)))) |
| 159 | if not optional and len(armnn_libs) == 0: |
| 160 | raise RuntimeError("""ArmNN library {} was not found in {}. Please install ArmNN to one of the standard |
| 161 | locations or set correct ARMNN_INCLUDE and ARMNN_LIB env variables.""".format(lib_name, |
| 162 | lib_search)) |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 163 | if optional and len(armnn_libs) == 0: |
| 164 | logger.warning("""Optional parser library %s was not found in %s and will not be installed.""", lib_name, |
| 165 | lib_search) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 166 | |
| 167 | # gives back tuple of names of the libs, set of unique libs locations and includes. |
| 168 | return list(armnn_libs.keys()), list(set( |
| 169 | map(lambda path: str(path.absolute().parent), armnn_libs.values()))) |
| 170 | |
| 171 | |
| 172 | class LazyArmnnFinderExtension(Extension): |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 173 | """Derived from `Extension` this class adds ArmNN libraries search on the user's machine. |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 174 | SWIG options and compilation flags are updated with relevant ArmNN libraries files locations (-L) and headers (-I). |
| 175 | |
| 176 | Search for ArmNN is executed only when attributes include_dirs, library_dirs, runtime_library_dirs, libraries or |
| 177 | swig_opts are queried. |
| 178 | |
| 179 | """ |
| 180 | |
| 181 | def __init__(self, name, sources, armnn_libs, include_dirs=None, define_macros=None, undef_macros=None, |
| 182 | library_dirs=None, |
| 183 | libraries=None, runtime_library_dirs=None, extra_objects=None, extra_compile_args=None, |
| 184 | extra_link_args=None, export_symbols=None, language=None, optional=None, **kw): |
| 185 | self._include_dirs = None |
| 186 | self._library_dirs = None |
| 187 | self._runtime_library_dirs = None |
| 188 | self._armnn_libs = armnn_libs |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 189 | self._optional = optional[0] |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 190 | # self.__swig_opts = None |
| 191 | super().__init__(name, sources, include_dirs, define_macros, undef_macros, library_dirs, libraries, |
| 192 | runtime_library_dirs, extra_objects, extra_compile_args, extra_link_args, export_symbols, |
| 193 | language, optional, **kw) |
| 194 | |
| 195 | @property |
| 196 | def include_dirs(self): |
| 197 | return self._include_dirs + find_includes() |
| 198 | |
| 199 | @include_dirs.setter |
| 200 | def include_dirs(self, include_dirs): |
| 201 | self._include_dirs = include_dirs |
| 202 | |
| 203 | @property |
| 204 | def library_dirs(self): |
| 205 | library_dirs = self._library_dirs |
| 206 | for lib in self._armnn_libs: |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 207 | _, lib_path = find_armnn(lib, self._optional) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 208 | library_dirs = library_dirs + lib_path |
| 209 | |
| 210 | return library_dirs |
| 211 | |
| 212 | @library_dirs.setter |
| 213 | def library_dirs(self, library_dirs): |
| 214 | self._library_dirs = library_dirs |
| 215 | |
| 216 | @property |
| 217 | def runtime_library_dirs(self): |
| 218 | library_dirs = self._runtime_library_dirs |
| 219 | for lib in self._armnn_libs: |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 220 | _, lib_path = find_armnn(lib, self._optional) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 221 | library_dirs = library_dirs + lib_path |
| 222 | |
| 223 | return library_dirs |
| 224 | |
| 225 | @runtime_library_dirs.setter |
| 226 | def runtime_library_dirs(self, runtime_library_dirs): |
| 227 | self._runtime_library_dirs = runtime_library_dirs |
| 228 | |
| 229 | @property |
| 230 | def libraries(self): |
| 231 | libraries = self._libraries |
| 232 | for lib in self._armnn_libs: |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 233 | lib_names, _ = find_armnn(lib, self._optional) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 234 | libraries = libraries + lib_names |
| 235 | |
| 236 | return libraries |
| 237 | |
| 238 | @libraries.setter |
| 239 | def libraries(self, libraries): |
| 240 | self._libraries = libraries |
| 241 | |
| 242 | def __eq__(self, other): |
| 243 | return self.__class__ == other.__class__ and self.name == other.name |
| 244 | |
| 245 | def __ne__(self, other): |
| 246 | return not self.__eq__(other) |
| 247 | |
| 248 | def __hash__(self): |
| 249 | return self.name.__hash__() |
| 250 | |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 251 | |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 252 | if __name__ == '__main__': |
| 253 | # mandatory extensions |
| 254 | pyarmnn_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn', |
| 255 | sources=['src/pyarmnn/_generated/armnn_wrap.cpp'], |
| 256 | extra_compile_args=['-std=c++14'], |
| 257 | language='c++', |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 258 | armnn_libs=['libarmnn.so'], |
| 259 | optional=[False] |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 260 | ) |
| 261 | pyarmnn_v_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn_version', |
| 262 | sources=['src/pyarmnn/_generated/armnn_version_wrap.cpp'], |
| 263 | extra_compile_args=['-std=c++14'], |
| 264 | language='c++', |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 265 | armnn_libs=['libarmnn.so'], |
| 266 | optional=[False] |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 267 | ) |
| 268 | extensions_to_build = [pyarmnn_v_module, pyarmnn_module] |
| 269 | |
| 270 | |
| 271 | # optional extensions |
| 272 | def add_parsers_ext(name: str, ext_list: list): |
| 273 | pyarmnn_optional_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn_{}'.format(name.lower()), |
| 274 | sources=['src/pyarmnn/_generated/armnn_{}_wrap.cpp'.format( |
| 275 | name.lower())], |
| 276 | extra_compile_args=['-std=c++14'], |
| 277 | language='c++', |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 278 | armnn_libs=['libarmnn.so', 'libarmnn{}.so'.format(name)], |
| 279 | optional=[True] |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 280 | ) |
| 281 | ext_list.append(pyarmnn_optional_module) |
| 282 | |
| 283 | |
| 284 | add_parsers_ext('CaffeParser', extensions_to_build) |
| 285 | add_parsers_ext('OnnxParser', extensions_to_build) |
| 286 | add_parsers_ext('TfParser', extensions_to_build) |
| 287 | add_parsers_ext('TfLiteParser', extensions_to_build) |
wangg | 5f960d9 | 2020-08-26 01:44:32 +0000 | [diff] [blame] | 288 | add_parsers_ext('Deserializer', extensions_to_build) |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 289 | |
| 290 | setup( |
| 291 | name='pyarmnn', |
| 292 | version=__version__, |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 293 | author='Arm Ltd, NXP Semiconductors', |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 294 | author_email='support@linaro.org', |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 295 | description=DOCLINES[0], |
| 296 | long_description="\n".join(DOCLINES[2:]), |
| 297 | url='https://mlplatform.org/', |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 298 | license='MIT', |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 299 | keywords='armnn neural network machine learning', |
| 300 | classifiers=[ |
| 301 | 'Development Status :: 3 - Alpha', |
| 302 | 'Intended Audience :: Developers', |
| 303 | 'Intended Audience :: Education', |
| 304 | 'Intended Audience :: Science/Research', |
| 305 | 'License :: OSI Approved :: MIT License', |
| 306 | 'Programming Language :: Python :: 3', |
| 307 | 'Programming Language :: Python :: 3 :: Only', |
| 308 | 'Programming Language :: Python :: 3.6', |
| 309 | 'Programming Language :: Python :: 3.7', |
| 310 | 'Programming Language :: Python :: 3.8', |
| 311 | 'Topic :: Scientific/Engineering', |
| 312 | 'Topic :: Scientific/Engineering :: Artificial Intelligence', |
| 313 | 'Topic :: Software Development', |
| 314 | 'Topic :: Software Development :: Libraries', |
| 315 | 'Topic :: Software Development :: Libraries :: Python Modules', |
| 316 | ], |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 317 | package_dir={'': 'src'}, |
| 318 | packages=[ |
| 319 | 'pyarmnn', |
| 320 | 'pyarmnn._generated', |
| 321 | 'pyarmnn._quantization', |
| 322 | 'pyarmnn._tensor', |
| 323 | 'pyarmnn._utilities' |
| 324 | ], |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 325 | data_files=[('', ['LICENSE'])], |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 326 | python_requires='>=3.5', |
| 327 | install_requires=['numpy'], |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 328 | cmdclass={ |
Éanna Ó Catháin | 6c3dee4 | 2020-09-10 13:02:37 +0100 | [diff] [blame] | 329 | 'build_py': ExtensionPriorityBuilder, |
Pavel Macenauer | 59e057f | 2020-04-15 14:17:26 +0000 | [diff] [blame] | 330 | 'build_ext': ArmnnVersionCheckerExtBuilder |
| 331 | }, |
Richard Burton | dc0c6ed | 2020-04-08 16:39:05 +0100 | [diff] [blame] | 332 | ext_modules=extensions_to_build |
| 333 | ) |