Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
302 lines (249 sloc) 10.1 KB
from distutils import dir_util, log
from distutils.command import build_ext
from distutils.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
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.c')
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, [])
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
You can’t perform that action at this time.