310 lines
10 KiB
Python
310 lines
10 KiB
Python
|
from setuptools import distutils as dutils
|
||
|
|
from setuptools.command import build_ext
|
||
|
|
from setuptools.extension import Extension
|
||
|
|
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
from numba.core import typing, sigutils
|
||
|
|
from numba.core.compiler_lock import global_compiler_lock
|
||
|
|
from numba.pycc.compiler import ModuleCompiler, ExportEntry
|
||
|
|
from numba.pycc.platform import Toolchain
|
||
|
|
from numba import cext
|
||
|
|
|
||
|
|
|
||
|
|
dir_util = dutils.dir_util
|
||
|
|
log = dutils.log
|
||
|
|
extension_libs = cext.get_extension_libs()
|
||
|
|
|
||
|
|
|
||
|
|
class CC(object):
|
||
|
|
"""
|
||
|
|
An ahead-of-time compiler to create extension modules that don't
|
||
|
|
depend on Numba.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# NOTE: using ccache can speed up repetitive builds
|
||
|
|
# (especially for the mixin modules)
|
||
|
|
|
||
|
|
_mixin_sources = ['modulemixin.c',] + extension_libs
|
||
|
|
|
||
|
|
# -flto strips all unused helper functions, which 1) makes the
|
||
|
|
# produced output much smaller and 2) can make the linking step faster.
|
||
|
|
# (the Windows linker seems to do this by default, judging by the results)
|
||
|
|
|
||
|
|
_extra_cflags = {
|
||
|
|
# Comment out due to odd behavior with GCC 4.9+ with LTO
|
||
|
|
# 'posix': ['-flto'],
|
||
|
|
}
|
||
|
|
|
||
|
|
_extra_ldflags = {
|
||
|
|
# Comment out due to odd behavior with GCC 4.9+ with LTO
|
||
|
|
# 'posix': ['-flto'],
|
||
|
|
}
|
||
|
|
|
||
|
|
def __init__(self, extension_name, source_module=None):
|
||
|
|
if '.' in extension_name:
|
||
|
|
raise ValueError("basename should be a simple module name, not "
|
||
|
|
"qualified name")
|
||
|
|
|
||
|
|
self._basename = extension_name
|
||
|
|
self._init_function = 'pycc_init_' + extension_name
|
||
|
|
self._exported_functions = {}
|
||
|
|
# Resolve source module name and directory
|
||
|
|
f = sys._getframe(1)
|
||
|
|
if source_module is None:
|
||
|
|
dct = f.f_globals
|
||
|
|
source_module = dct['__name__']
|
||
|
|
elif hasattr(source_module, '__name__'):
|
||
|
|
dct = source_module.__dict__
|
||
|
|
source_module = source_module.__name__
|
||
|
|
else:
|
||
|
|
dct = sys.modules[source_module].__dict__
|
||
|
|
|
||
|
|
self._source_path = dct.get('__file__', '')
|
||
|
|
self._source_module = source_module
|
||
|
|
self._toolchain = Toolchain()
|
||
|
|
self._verbose = False
|
||
|
|
# By default, output in directory of caller module
|
||
|
|
self._output_dir = os.path.dirname(self._source_path)
|
||
|
|
self._output_file = self._toolchain.get_ext_filename(extension_name)
|
||
|
|
self._use_nrt = True
|
||
|
|
self._target_cpu = ''
|
||
|
|
|
||
|
|
@property
|
||
|
|
def name(self):
|
||
|
|
"""
|
||
|
|
The name of the extension module to create.
|
||
|
|
"""
|
||
|
|
return self._basename
|
||
|
|
|
||
|
|
@property
|
||
|
|
def output_file(self):
|
||
|
|
"""
|
||
|
|
The specific output file (a DLL) that will be generated.
|
||
|
|
"""
|
||
|
|
return self._output_file
|
||
|
|
|
||
|
|
@output_file.setter
|
||
|
|
def output_file(self, value):
|
||
|
|
self._output_file = value
|
||
|
|
|
||
|
|
@property
|
||
|
|
def output_dir(self):
|
||
|
|
"""
|
||
|
|
The directory the output file will be put in.
|
||
|
|
"""
|
||
|
|
return self._output_dir
|
||
|
|
|
||
|
|
@output_dir.setter
|
||
|
|
def output_dir(self, value):
|
||
|
|
self._output_dir = value
|
||
|
|
|
||
|
|
@property
|
||
|
|
def use_nrt(self):
|
||
|
|
return self._use_nrt
|
||
|
|
|
||
|
|
@use_nrt.setter
|
||
|
|
def use_nrt(self, value):
|
||
|
|
self._use_nrt = value
|
||
|
|
|
||
|
|
@property
|
||
|
|
def target_cpu(self):
|
||
|
|
"""
|
||
|
|
The target CPU model for code generation.
|
||
|
|
"""
|
||
|
|
return self._target_cpu
|
||
|
|
|
||
|
|
@target_cpu.setter
|
||
|
|
def target_cpu(self, value):
|
||
|
|
self._target_cpu = value
|
||
|
|
|
||
|
|
@property
|
||
|
|
def verbose(self):
|
||
|
|
"""
|
||
|
|
Whether to display detailed information when compiling.
|
||
|
|
"""
|
||
|
|
return self._verbose
|
||
|
|
|
||
|
|
@verbose.setter
|
||
|
|
def verbose(self, value):
|
||
|
|
self._verbose = value
|
||
|
|
|
||
|
|
def export(self, exported_name, sig):
|
||
|
|
"""
|
||
|
|
Mark a function for exporting in the extension module.
|
||
|
|
"""
|
||
|
|
fn_args, fn_retty = sigutils.normalize_signature(sig)
|
||
|
|
sig = typing.signature(fn_retty, *fn_args)
|
||
|
|
if exported_name in self._exported_functions:
|
||
|
|
raise KeyError("duplicated export symbol %s" % (exported_name))
|
||
|
|
|
||
|
|
def decorator(func):
|
||
|
|
entry = ExportEntry(exported_name, sig, func)
|
||
|
|
self._exported_functions[exported_name] = entry
|
||
|
|
return func
|
||
|
|
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _export_entries(self):
|
||
|
|
return sorted(self._exported_functions.values(),
|
||
|
|
key=lambda entry: entry.symbol)
|
||
|
|
|
||
|
|
def _get_mixin_sources(self):
|
||
|
|
here = os.path.dirname(__file__)
|
||
|
|
mixin_sources = self._mixin_sources[:]
|
||
|
|
if self._use_nrt:
|
||
|
|
mixin_sources.append('../core/runtime/nrt.cpp')
|
||
|
|
return [os.path.join(here, f) for f in mixin_sources]
|
||
|
|
|
||
|
|
def _get_mixin_defines(self):
|
||
|
|
# Macro definitions required by modulemixin.c
|
||
|
|
return [
|
||
|
|
('PYCC_MODULE_NAME', self._basename),
|
||
|
|
('PYCC_USE_NRT', int(self._use_nrt)),
|
||
|
|
]
|
||
|
|
|
||
|
|
def _get_extra_cflags(self):
|
||
|
|
extra_cflags = self._extra_cflags.get(sys.platform, [])
|
||
|
|
if not extra_cflags:
|
||
|
|
extra_cflags = self._extra_cflags.get(os.name, [])
|
||
|
|
return extra_cflags
|
||
|
|
|
||
|
|
def _get_extra_ldflags(self):
|
||
|
|
extra_ldflags = self._extra_ldflags.get(sys.platform, [])
|
||
|
|
if not extra_ldflags:
|
||
|
|
extra_ldflags = self._extra_ldflags.get(os.name, [])
|
||
|
|
# helperlib uses pthread on linux. make sure we are linking to it.
|
||
|
|
if sys.platform.startswith("linux"):
|
||
|
|
if "-pthread" not in extra_ldflags:
|
||
|
|
extra_ldflags.append('-pthread')
|
||
|
|
return extra_ldflags
|
||
|
|
|
||
|
|
def _compile_mixins(self, build_dir):
|
||
|
|
sources = self._get_mixin_sources()
|
||
|
|
macros = self._get_mixin_defines()
|
||
|
|
include_dirs = self._toolchain.get_python_include_dirs()
|
||
|
|
|
||
|
|
extra_cflags = self._get_extra_cflags()
|
||
|
|
# XXX distutils creates a whole subtree inside build_dir,
|
||
|
|
# e.g. /tmp/test_pycc/home/antoine/numba/numba/pycc/modulemixin.o
|
||
|
|
objects = self._toolchain.compile_objects(sources, build_dir,
|
||
|
|
include_dirs=include_dirs,
|
||
|
|
macros=macros,
|
||
|
|
extra_cflags=extra_cflags)
|
||
|
|
return objects
|
||
|
|
|
||
|
|
@global_compiler_lock
|
||
|
|
def _compile_object_files(self, build_dir):
|
||
|
|
compiler = ModuleCompiler(self._export_entries, self._basename,
|
||
|
|
self._use_nrt, cpu_name=self._target_cpu)
|
||
|
|
compiler.external_init_function = self._init_function
|
||
|
|
temp_obj = os.path.join(build_dir,
|
||
|
|
os.path.splitext(self._output_file)[0] + '.o')
|
||
|
|
log.info("generating LLVM code for '%s' into %s",
|
||
|
|
self._basename, temp_obj)
|
||
|
|
compiler.write_native_object(temp_obj, wrap=True)
|
||
|
|
return [temp_obj], compiler.dll_exports
|
||
|
|
|
||
|
|
@global_compiler_lock
|
||
|
|
def compile(self):
|
||
|
|
"""
|
||
|
|
Compile the extension module.
|
||
|
|
"""
|
||
|
|
self._toolchain.verbose = self.verbose
|
||
|
|
build_dir = tempfile.mkdtemp(prefix='pycc-build-%s-' % self._basename)
|
||
|
|
|
||
|
|
# Compile object file
|
||
|
|
objects, dll_exports = self._compile_object_files(build_dir)
|
||
|
|
|
||
|
|
# Compile mixins
|
||
|
|
objects += self._compile_mixins(build_dir)
|
||
|
|
|
||
|
|
# Then create shared library
|
||
|
|
extra_ldflags = self._get_extra_ldflags()
|
||
|
|
output_dll = os.path.join(self._output_dir, self._output_file)
|
||
|
|
libraries = self._toolchain.get_python_libraries()
|
||
|
|
library_dirs = self._toolchain.get_python_library_dirs()
|
||
|
|
self._toolchain.link_shared(output_dll, objects,
|
||
|
|
libraries, library_dirs,
|
||
|
|
export_symbols=dll_exports,
|
||
|
|
extra_ldflags=extra_ldflags)
|
||
|
|
|
||
|
|
shutil.rmtree(build_dir)
|
||
|
|
|
||
|
|
def distutils_extension(self, **kwargs):
|
||
|
|
"""
|
||
|
|
Create a distutils extension object that can be used in your
|
||
|
|
setup.py.
|
||
|
|
"""
|
||
|
|
macros = kwargs.pop('macros', []) + self._get_mixin_defines()
|
||
|
|
depends = kwargs.pop('depends', []) + [self._source_path]
|
||
|
|
extra_compile_args = (kwargs.pop('extra_compile_args', [])
|
||
|
|
+ self._get_extra_cflags())
|
||
|
|
extra_link_args = (kwargs.pop('extra_link_args', [])
|
||
|
|
+ self._get_extra_ldflags())
|
||
|
|
include_dirs = (kwargs.pop('include_dirs', [])
|
||
|
|
+ self._toolchain.get_python_include_dirs())
|
||
|
|
libraries = (kwargs.pop('libraries', [])
|
||
|
|
+ self._toolchain.get_python_libraries())
|
||
|
|
library_dirs = (kwargs.pop('library_dirs', [])
|
||
|
|
+ self._toolchain.get_python_library_dirs())
|
||
|
|
python_package_path = self._source_module[:self._source_module.rfind('.')+1]
|
||
|
|
|
||
|
|
ext = _CCExtension(name=python_package_path + self._basename,
|
||
|
|
sources=self._get_mixin_sources(),
|
||
|
|
depends=depends,
|
||
|
|
define_macros=macros,
|
||
|
|
include_dirs=include_dirs,
|
||
|
|
libraries=libraries,
|
||
|
|
library_dirs=library_dirs,
|
||
|
|
extra_compile_args=extra_compile_args,
|
||
|
|
extra_link_args=extra_link_args,
|
||
|
|
**kwargs)
|
||
|
|
ext.monkey_patch_distutils()
|
||
|
|
ext._cc = self
|
||
|
|
return ext
|
||
|
|
|
||
|
|
|
||
|
|
class _CCExtension(Extension):
|
||
|
|
"""
|
||
|
|
A Numba-specific Extension subclass to LLVM-compile pure Python code
|
||
|
|
to an extension module.
|
||
|
|
"""
|
||
|
|
|
||
|
|
_cc = None
|
||
|
|
_distutils_monkey_patched = False
|
||
|
|
|
||
|
|
def _prepare_object_files(self, build_ext):
|
||
|
|
cc = self._cc
|
||
|
|
dir_util.mkpath(os.path.join(build_ext.build_temp, *self.name.split('.')[:-1]))
|
||
|
|
objects, _ = cc._compile_object_files(build_ext.build_temp)
|
||
|
|
# Add generated object files for linking
|
||
|
|
self.extra_objects = objects
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def monkey_patch_distutils(cls):
|
||
|
|
"""
|
||
|
|
Monkey-patch distutils with our own build_ext class knowing
|
||
|
|
about pycc-compiled extensions modules.
|
||
|
|
"""
|
||
|
|
if cls._distutils_monkey_patched:
|
||
|
|
return
|
||
|
|
|
||
|
|
_orig_build_ext = build_ext.build_ext
|
||
|
|
|
||
|
|
class _CC_build_ext(_orig_build_ext):
|
||
|
|
|
||
|
|
def build_extension(self, ext):
|
||
|
|
if isinstance(ext, _CCExtension):
|
||
|
|
ext._prepare_object_files(self)
|
||
|
|
|
||
|
|
_orig_build_ext.build_extension(self, ext)
|
||
|
|
|
||
|
|
build_ext.build_ext = _CC_build_ext
|
||
|
|
|
||
|
|
cls._distutils_monkey_patched = True
|