CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
cc.py310 linesDownload Raw Back to pycc
1from setuptools import distutils as dutils
2from setuptools.command import build_ext
3from setuptools.extension import Extension
4
5import os
6import shutil
7import sys
8import tempfile
9
10from numba.core import typing, sigutils
11from numba.core.compiler_lock import global_compiler_lock
12from numba.pycc.compiler import ModuleCompiler, ExportEntry
13from numba.pycc.platform import Toolchain
14from numba import cext
15
16
17dir_util = dutils.dir_util
18log = dutils.log
19extension_libs = cext.get_extension_libs()
20
21
22class CC(object):
23    """
24    An ahead-of-time compiler to create extension modules that don't
25    depend on Numba.
26    """
27
28    # NOTE: using ccache can speed up repetitive builds
29    # (especially for the mixin modules)
30
31    _mixin_sources = ['modulemixin.c',]  + extension_libs
32
33    # -flto strips all unused helper functions, which 1) makes the
34    # produced output much smaller and 2) can make the linking step faster.
35    # (the Windows linker seems to do this by default, judging by the results)
36
37    _extra_cflags = {
38        # Comment out due to odd behavior with GCC 4.9+ with LTO
39        # 'posix': ['-flto'],
40        }
41
42    _extra_ldflags = {
43        # Comment out due to odd behavior with GCC 4.9+ with LTO
44        # 'posix': ['-flto'],
45        }
46
47    def __init__(self, extension_name, source_module=None):
48        if '.' in extension_name:
49            raise ValueError("basename should be a simple module name, not "
50                             "qualified name")
51
52        self._basename = extension_name
53        self._init_function = 'pycc_init_' + extension_name
54        self._exported_functions = {}
55        # Resolve source module name and directory
56        f = sys._getframe(1)
57        if source_module is None:
58            dct = f.f_globals
59            source_module = dct['__name__']
60        elif hasattr(source_module, '__name__'):
61            dct = source_module.__dict__
62            source_module = source_module.__name__
63        else:
64            dct = sys.modules[source_module].__dict__
65
66        self._source_path = dct.get('__file__', '')
67        self._source_module = source_module
68        self._toolchain = Toolchain()
69        self._verbose = False
70        # By default, output in directory of caller module
71        self._output_dir = os.path.dirname(self._source_path)
72        self._output_file = self._toolchain.get_ext_filename(extension_name)
73        self._use_nrt = True
74        self._target_cpu = ''
75
76    @property
77    def name(self):
78        """
79        The name of the extension module to create.
80        """
81        return self._basename
82
83    @property
84    def output_file(self):
85        """
86        The specific output file (a DLL) that will be generated.
87        """
88        return self._output_file
89
90    @output_file.setter
91    def output_file(self, value):
92        self._output_file = value
93
94    @property
95    def output_dir(self):
96        """
97        The directory the output file will be put in.
98        """
99        return self._output_dir
100
101    @output_dir.setter
102    def output_dir(self, value):
103        self._output_dir = value
104
105    @property
106    def use_nrt(self):
107        return self._use_nrt
108
109    @use_nrt.setter
110    def use_nrt(self, value):
111        self._use_nrt = value
112
113    @property
114    def target_cpu(self):
115        """
116        The target CPU model for code generation.
117        """
118        return self._target_cpu
119
120    @target_cpu.setter
121    def target_cpu(self, value):
122        self._target_cpu = value
123
124    @property
125    def verbose(self):
126        """
127        Whether to display detailed information when compiling.
128        """
129        return self._verbose
130
131    @verbose.setter
132    def verbose(self, value):
133        self._verbose = value
134
135    def export(self, exported_name, sig):
136        """
137        Mark a function for exporting in the extension module.
138        """
139        fn_args, fn_retty = sigutils.normalize_signature(sig)
140        sig = typing.signature(fn_retty, *fn_args)
141        if exported_name in self._exported_functions:
142            raise KeyError("duplicated export symbol %s" % (exported_name))
143
144        def decorator(func):
145            entry = ExportEntry(exported_name, sig, func)
146            self._exported_functions[exported_name] = entry
147            return func
148
149        return decorator
150
151    @property
152    def _export_entries(self):
153        return sorted(self._exported_functions.values(),
154                      key=lambda entry: entry.symbol)
155
156    def _get_mixin_sources(self):
157        here = os.path.dirname(__file__)
158        mixin_sources = self._mixin_sources[:]
159        if self._use_nrt:
160            mixin_sources.append('../core/runtime/nrt.cpp')
161        return [os.path.join(here, f) for f in mixin_sources]
162
163    def _get_mixin_defines(self):
164        # Macro definitions required by modulemixin.c
165        return [
166            ('PYCC_MODULE_NAME', self._basename),
167            ('PYCC_USE_NRT', int(self._use_nrt)),
168            ]
169
170    def _get_extra_cflags(self):
171        extra_cflags = self._extra_cflags.get(sys.platform, [])
172        if not extra_cflags:
173            extra_cflags = self._extra_cflags.get(os.name, [])
174        return extra_cflags
175
176    def _get_extra_ldflags(self):
177        extra_ldflags = self._extra_ldflags.get(sys.platform, [])
178        if not extra_ldflags:
179            extra_ldflags = self._extra_ldflags.get(os.name, [])
180        # helperlib uses pthread on linux. make sure we are linking to it.
181        if sys.platform.startswith("linux"):
182            if "-pthread" not in extra_ldflags:
183                extra_ldflags.append('-pthread')
184        return extra_ldflags
185
186    def _compile_mixins(self, build_dir):
187        sources = self._get_mixin_sources()
188        macros = self._get_mixin_defines()
189        include_dirs = self._toolchain.get_python_include_dirs()
190
191        extra_cflags = self._get_extra_cflags()
192        # XXX distutils creates a whole subtree inside build_dir,
193        # e.g. /tmp/test_pycc/home/antoine/numba/numba/pycc/modulemixin.o
194        objects = self._toolchain.compile_objects(sources, build_dir,
195                                                  include_dirs=include_dirs,
196                                                  macros=macros,
197                                                  extra_cflags=extra_cflags)
198        return objects
199
200    @global_compiler_lock
201    def _compile_object_files(self, build_dir):
202        compiler = ModuleCompiler(self._export_entries, self._basename,
203                                self._use_nrt, cpu_name=self._target_cpu)
204        compiler.external_init_function = self._init_function
205        temp_obj = os.path.join(build_dir,
206                                os.path.splitext(self._output_file)[0] + '.o')
207        log.info("generating LLVM code for '%s' into %s",
208                self._basename, temp_obj)
209        compiler.write_native_object(temp_obj, wrap=True)
210        return [temp_obj], compiler.dll_exports
211
212    @global_compiler_lock
213    def compile(self):
214        """
215        Compile the extension module.
216        """
217        self._toolchain.verbose = self.verbose
218        build_dir = tempfile.mkdtemp(prefix='pycc-build-%s-' % self._basename)
219
220        # Compile object file
221        objects, dll_exports = self._compile_object_files(build_dir)
222
223        # Compile mixins
224        objects += self._compile_mixins(build_dir)
225
226        # Then create shared library
227        extra_ldflags = self._get_extra_ldflags()
228        output_dll = os.path.join(self._output_dir, self._output_file)
229        libraries = self._toolchain.get_python_libraries()
230        library_dirs = self._toolchain.get_python_library_dirs()
231        self._toolchain.link_shared(output_dll, objects,
232                                    libraries, library_dirs,
233                                    export_symbols=dll_exports,
234                                    extra_ldflags=extra_ldflags)
235
236        shutil.rmtree(build_dir)
237
238    def distutils_extension(self, **kwargs):
239        """
240        Create a distutils extension object that can be used in your
241        setup.py.
242        """
243        macros = kwargs.pop('macros', []) + self._get_mixin_defines()
244        depends = kwargs.pop('depends', []) + [self._source_path]
245        extra_compile_args = (kwargs.pop('extra_compile_args', [])
246                              + self._get_extra_cflags())
247        extra_link_args = (kwargs.pop('extra_link_args', [])
248                           + self._get_extra_ldflags())
249        include_dirs = (kwargs.pop('include_dirs', [])
250                        + self._toolchain.get_python_include_dirs())
251        libraries = (kwargs.pop('libraries', [])
252                     + self._toolchain.get_python_libraries())
253        library_dirs = (kwargs.pop('library_dirs', [])
254                        + self._toolchain.get_python_library_dirs())
255        python_package_path = self._source_module[:self._source_module.rfind('.')+1]
256
257        ext = _CCExtension(name=python_package_path + self._basename,
258                           sources=self._get_mixin_sources(),
259                           depends=depends,
260                           define_macros=macros,
261                           include_dirs=include_dirs,
262                           libraries=libraries,
263                           library_dirs=library_dirs,
264                           extra_compile_args=extra_compile_args,
265                           extra_link_args=extra_link_args,
266                           **kwargs)
267        ext.monkey_patch_distutils()
268        ext._cc = self
269        return ext
270
271
272class _CCExtension(Extension):
273    """
274    A Numba-specific Extension subclass to LLVM-compile pure Python code
275    to an extension module.
276    """
277
278    _cc = None
279    _distutils_monkey_patched = False
280
281    def _prepare_object_files(self, build_ext):
282        cc = self._cc
283        dir_util.mkpath(os.path.join(build_ext.build_temp, *self.name.split('.')[:-1]))
284        objects, _ = cc._compile_object_files(build_ext.build_temp)
285        # Add generated object files for linking
286        self.extra_objects = objects
287
288    @classmethod
289    def monkey_patch_distutils(cls):
290        """
291        Monkey-patch distutils with our own build_ext class knowing
292        about pycc-compiled extensions modules.
293        """
294        if cls._distutils_monkey_patched:
295            return
296
297        _orig_build_ext = build_ext.build_ext
298
299        class _CC_build_ext(_orig_build_ext):
300
301            def build_extension(self, ext):
302                if isinstance(ext, _CCExtension):
303                    ext._prepare_object_files(self)
304
305                _orig_build_ext.build_extension(self, ext)
306
307        build_ext.build_ext = _CC_build_ext
308
309        cls._distutils_monkey_patched = True
310 
Aluode/PerceptionLabPortable · CoolFace