Kshitij Sisodia | f9efe0d | 2022-09-30 16:42:50 +0100 | [diff] [blame] | 1 | # |
| 2 | # SPDX-FileCopyrightText: Copyright 2021-2022 Arm Limited and/or its affiliates <open-source-office@arm.com> |
| 3 | # SPDX-License-Identifier: Apache-2.0 |
| 4 | # |
| 5 | import logging |
| 6 | import os |
| 7 | from functools import lru_cache |
| 8 | from itertools import chain |
| 9 | from pathlib import Path |
| 10 | |
| 11 | from setuptools import setup |
| 12 | from distutils.core import Extension |
| 13 | from setuptools.command.build_py import build_py |
| 14 | |
| 15 | logger = logging.Logger(__name__) |
| 16 | |
| 17 | |
| 18 | def linux_gcc_lib_search(): |
| 19 | """ |
| 20 | Calls the `gcc` to get linker default system paths. |
| 21 | Returns: |
| 22 | list of paths |
| 23 | """ |
| 24 | cmd = 'gcc --print-search-dirs | grep libraries' |
| 25 | cmd_res = os.popen(cmd).read() |
| 26 | cmd_res = cmd_res.split('=') |
| 27 | if len(cmd_res) > 1: |
| 28 | return tuple(cmd_res[1].split(':')) |
| 29 | return None |
| 30 | |
| 31 | |
| 32 | def find_includes(include_env: str = 'ETHOS_U_DRIVER_INCLUDE'): |
| 33 | include_path = os.getenv(include_env, '') |
| 34 | return [include_path] if include_path else ['/usr/local/include', '/usr/include'] |
| 35 | |
| 36 | |
| 37 | @lru_cache(maxsize=1) |
| 38 | def find_driver(lib_name: str, |
| 39 | optional: bool = False, |
| 40 | libs_env: str = 'ETHOS_U_DRIVER_LIB', |
| 41 | default_lib_search: tuple = linux_gcc_lib_search()): |
| 42 | """ |
| 43 | Searches for library installation on the local machine. |
| 44 | |
| 45 | Args: |
| 46 | lib_name: library name to find |
| 47 | optional: Do not fail if optional. Default is False - fail if library was not found. |
| 48 | libs_env: custom environment variable pointing to libraries location, default is 'ETHOS_U_DRIVER_LIB' |
| 49 | default_lib_search: list of paths to search for a library if not found within path provided by |
| 50 | 'ETHOS_U_DRIVER_LIB' env variable |
| 51 | |
| 52 | Returns: |
| 53 | tuple containing name of the driver libs, paths to the libs |
| 54 | """ |
| 55 | |
| 56 | lib_path = os.getenv(libs_env, "") |
| 57 | |
| 58 | lib_search = [lib_path] if lib_path else default_lib_search |
| 59 | |
| 60 | libs = dict(map(lambda path: (':{}'.format(path.name), path), |
| 61 | chain.from_iterable(map(lambda lib_path: Path(lib_path).glob(lib_name), |
| 62 | lib_search)))) |
| 63 | if not optional and len(libs) == 0: |
| 64 | raise RuntimeError("""Ethos-U driver library {} was not found in {}. Please install driver to one of the standard |
| 65 | locations or set correct ETHOS_U_DRIVER_INCLUDE and ETHOS_U_DRIVER_LIB env variables.""" |
| 66 | .format(lib_name, lib_search)) |
| 67 | |
| 68 | # gives back tuple of names of the libs, set of unique libs locations and includes. |
| 69 | return list(libs.keys()), list(set( |
| 70 | map(lambda path: str(path.absolute().parent), libs.values()))) |
| 71 | |
| 72 | |
| 73 | class LibFinderExtension(Extension): |
| 74 | """ |
| 75 | Derived from `Extension` this class adds libraries search on the user's machine. |
| 76 | SWIG options and compilation flags are updated with relevant libraries files locations (-L) and headers (-I). |
| 77 | |
| 78 | Search for the library is executed only when attributes include_dirs, library_dirs, runtime_library_dirs, libraries or |
| 79 | swig_opts are queried. |
| 80 | |
| 81 | """ |
| 82 | |
| 83 | def __init__(self, name, sources, libs, include_dirs=None, define_macros=None, undef_macros=None, |
| 84 | library_dirs=None, |
| 85 | libraries=None, runtime_library_dirs=None, extra_objects=None, extra_compile_args=None, |
| 86 | extra_link_args=None, export_symbols=None, language=None, **kw): |
| 87 | self._include_dirs = None |
| 88 | self._library_dirs = None |
| 89 | self._runtime_library_dirs = None |
| 90 | self._libs = libs |
| 91 | super().__init__(name, sources, include_dirs, define_macros, undef_macros, library_dirs, libraries, |
| 92 | runtime_library_dirs, extra_objects, extra_compile_args, extra_link_args, export_symbols, |
| 93 | language, **kw) |
| 94 | |
| 95 | @property |
| 96 | def include_dirs(self): |
| 97 | return self._include_dirs + find_includes() |
| 98 | |
| 99 | @include_dirs.setter |
| 100 | def include_dirs(self, include_dirs): |
| 101 | self._include_dirs = include_dirs |
| 102 | |
| 103 | @property |
| 104 | def library_dirs(self): |
| 105 | library_dirs = self._library_dirs |
| 106 | for lib in self._libs: |
| 107 | _, lib_path = find_driver(lib) |
| 108 | library_dirs = library_dirs + lib_path |
| 109 | |
| 110 | return library_dirs |
| 111 | |
| 112 | @library_dirs.setter |
| 113 | def library_dirs(self, library_dirs): |
| 114 | self._library_dirs = library_dirs |
| 115 | |
| 116 | @property |
| 117 | def runtime_library_dirs(self): |
| 118 | library_dirs = self._runtime_library_dirs |
| 119 | for lib in self._libs: |
| 120 | _, lib_path = find_driver(lib) |
| 121 | library_dirs = library_dirs + lib_path |
| 122 | |
| 123 | return library_dirs |
| 124 | |
| 125 | @runtime_library_dirs.setter |
| 126 | def runtime_library_dirs(self, runtime_library_dirs): |
| 127 | self._runtime_library_dirs = runtime_library_dirs |
| 128 | |
| 129 | @property |
| 130 | def libraries(self): |
| 131 | libraries = self._libraries |
| 132 | for lib in self._libs: |
| 133 | lib_names, _ = find_driver(lib) |
| 134 | libraries = libraries + lib_names |
| 135 | |
| 136 | return libraries |
| 137 | |
| 138 | @libraries.setter |
| 139 | def libraries(self, libraries): |
| 140 | self._libraries = libraries |
| 141 | |
| 142 | def __eq__(self, other): |
| 143 | return self.__class__ == other.__class__ and self.name == other.name |
| 144 | |
| 145 | def __ne__(self, other): |
| 146 | return not self.__eq__(other) |
| 147 | |
| 148 | def __hash__(self): |
| 149 | return self.name.__hash__() |
| 150 | |
| 151 | |
| 152 | class ExtensionPriorityBuilder(build_py): |
| 153 | """ |
| 154 | Runs extension builder before other stages. Otherwise generated files are not included to the distribution. |
| 155 | """ |
| 156 | |
| 157 | def run(self): |
| 158 | self.run_command('build_ext') |
| 159 | return super().run() |
| 160 | |
| 161 | |
| 162 | if __name__ == '__main__': |
| 163 | # mandatory extensions |
| 164 | driver_module = LibFinderExtension('ethosu_driver._generated._driver', |
| 165 | sources=['src/ethosu_driver/_generated/driver_wrap.cpp'], |
| 166 | extra_compile_args=['-std=gnu++14'], |
| 167 | language='c++', |
| 168 | libs=['libethosu.a'] |
| 169 | ) |
| 170 | |
| 171 | extensions_to_build = [driver_module] |
| 172 | |
| 173 | setup( |
| 174 | name='ethosu_driver', |
| 175 | version='1.0.0', |
| 176 | author='Arm ltd', |
| 177 | author_email='support@arm.com', |
| 178 | description='Arm Ethos-U NPU Linux Driver Stack Python wrapper', |
| 179 | url='https://git.mlplatform.org/ml/ethos-u/ethos-u-linux-driver-stack.git/', |
| 180 | license='Apache License 2.0', |
| 181 | classifiers=[ |
| 182 | "Development Status :: 5 - Production/Stable", |
| 183 | "Intended Audience :: Developers", |
| 184 | "License :: OSI Approved :: Apache Software License", |
| 185 | "Operating System :: POSIX :: Linux", |
| 186 | "Programming Language :: C", |
| 187 | "Programming Language :: Python :: 3", |
| 188 | "Programming Language :: Python :: 3.5", |
| 189 | "Topic :: Scientific/Engineering :: Artificial Intelligence" |
| 190 | ], |
| 191 | keywords=["ethos-u", "driver", "npu"], |
| 192 | package_dir={'': 'src'}, |
| 193 | packages=[ |
| 194 | 'ethosu_driver', |
| 195 | 'ethosu_driver._generated', |
| 196 | 'ethosu_driver._utilities' |
| 197 | ], |
| 198 | data_files=[('', ['LICENSE'])], |
| 199 | entry_points={"console_scripts": ["inference_runner = ethosu_driver.inference_runner:main"]}, |
| 200 | python_requires='>=3.5', |
| 201 | extras_require={"numpy": ["numpy"]}, |
| 202 | cmdclass={'build_py': ExtensionPriorityBuilder}, |
| 203 | ext_modules=extensions_to_build |
| 204 | ) |