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