478 lines
18 KiB
Python
478 lines
18 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from llvmlite import ir
|
||
|
|
from llvmlite.binding import Linkage
|
||
|
|
|
||
|
|
from numba.pycc import llvm_types as lt
|
||
|
|
from numba.core.cgutils import create_constant_array
|
||
|
|
from numba.core.compiler import compile_extra, Flags
|
||
|
|
from numba.core.compiler_lock import global_compiler_lock
|
||
|
|
|
||
|
|
from numba.core.registry import cpu_target
|
||
|
|
from numba.core.runtime import nrtdynmod
|
||
|
|
from numba.core import cgutils
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
__all__ = ['Compiler']
|
||
|
|
|
||
|
|
NULL = ir.Constant(lt._void_star, None)
|
||
|
|
ZERO = ir.Constant(lt._int32, 0)
|
||
|
|
ONE = ir.Constant(lt._int32, 1)
|
||
|
|
METH_VARARGS_AND_KEYWORDS = ir.Constant(lt._int32, 1|2)
|
||
|
|
|
||
|
|
|
||
|
|
def get_header():
|
||
|
|
import numpy
|
||
|
|
import textwrap
|
||
|
|
|
||
|
|
return textwrap.dedent("""\
|
||
|
|
#include <stdint.h>
|
||
|
|
|
||
|
|
#ifndef HAVE_LONGDOUBLE
|
||
|
|
#define HAVE_LONGDOUBLE %d
|
||
|
|
#endif
|
||
|
|
|
||
|
|
typedef struct {
|
||
|
|
float real;
|
||
|
|
float imag;
|
||
|
|
} complex64;
|
||
|
|
|
||
|
|
typedef struct {
|
||
|
|
double real;
|
||
|
|
double imag;
|
||
|
|
} complex128;
|
||
|
|
|
||
|
|
#if HAVE_LONGDOUBLE
|
||
|
|
typedef struct {
|
||
|
|
long double real;
|
||
|
|
long double imag;
|
||
|
|
} complex256;
|
||
|
|
#endif
|
||
|
|
|
||
|
|
typedef float float32;
|
||
|
|
typedef double float64;
|
||
|
|
#if HAVE_LONGDOUBLE
|
||
|
|
typedef long double float128;
|
||
|
|
#endif
|
||
|
|
""" % hasattr(numpy, 'complex256'))
|
||
|
|
|
||
|
|
|
||
|
|
class ExportEntry(object):
|
||
|
|
"""
|
||
|
|
A simple record for exporting symbols.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, symbol, signature, function):
|
||
|
|
self.symbol = symbol
|
||
|
|
self.signature = signature
|
||
|
|
self.function = function
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
return "ExportEntry(%r, %r)" % (self.symbol, self.signature)
|
||
|
|
|
||
|
|
|
||
|
|
class _ModuleCompiler(object):
|
||
|
|
"""A base class to compile Python modules to a single shared library or
|
||
|
|
extension module.
|
||
|
|
|
||
|
|
:param export_entries: a list of ExportEntry instances.
|
||
|
|
:param module_name: the name of the exported module.
|
||
|
|
"""
|
||
|
|
|
||
|
|
#: Structure used to describe a method of an extension type.
|
||
|
|
#: struct PyMethodDef {
|
||
|
|
#: const char *ml_name; /* The name of the built-in function/method */
|
||
|
|
#: PyCFunction ml_meth; /* The C function that implements it */
|
||
|
|
#: int ml_flags; /* Combination of METH_xxx flags, which mostly
|
||
|
|
#: describe the args expected by the C func */
|
||
|
|
#: const char *ml_doc; /* The __doc__ attribute, or NULL */
|
||
|
|
#: };
|
||
|
|
method_def_ty = ir.LiteralStructType((lt._int8_star,
|
||
|
|
lt._void_star,
|
||
|
|
lt._int32,
|
||
|
|
lt._int8_star))
|
||
|
|
|
||
|
|
method_def_ptr = ir.PointerType(method_def_ty)
|
||
|
|
# The structure type constructed by PythonAPI.serialize_uncached()
|
||
|
|
# when updating this, also make sure to update `env_def_t` struct in
|
||
|
|
# numba/pycc/modulemixin.c
|
||
|
|
env_def_ty = ir.LiteralStructType((lt._void_star,
|
||
|
|
lt._int32,
|
||
|
|
lt._void_star,
|
||
|
|
lt._void_star,
|
||
|
|
lt._int32))
|
||
|
|
env_def_ptr = ir.PointerType(env_def_ty)
|
||
|
|
|
||
|
|
def __init__(self, export_entries, module_name, use_nrt=False,
|
||
|
|
**aot_options):
|
||
|
|
self.module_name = module_name
|
||
|
|
self.export_python_wrap = False
|
||
|
|
self.dll_exports = []
|
||
|
|
self.export_entries = export_entries
|
||
|
|
# Used by the CC API but not the legacy API
|
||
|
|
self.external_init_function = None
|
||
|
|
self.use_nrt = use_nrt
|
||
|
|
|
||
|
|
self.typing_context = cpu_target.typing_context
|
||
|
|
self.context = cpu_target.target_context.with_aot_codegen(
|
||
|
|
self.module_name, **aot_options)
|
||
|
|
|
||
|
|
def _mangle_method_symbol(self, func_name):
|
||
|
|
return "._pycc_method_%s" % (func_name,)
|
||
|
|
|
||
|
|
def _emit_python_wrapper(self, llvm_module):
|
||
|
|
"""Emit generated Python wrapper and extension module code.
|
||
|
|
"""
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
@global_compiler_lock
|
||
|
|
def _cull_exports(self):
|
||
|
|
"""Read all the exported functions/modules in the translator
|
||
|
|
environment, and join them into a single LLVM module.
|
||
|
|
"""
|
||
|
|
self.exported_function_types = {}
|
||
|
|
self.function_environments = {}
|
||
|
|
self.environment_gvs = {}
|
||
|
|
|
||
|
|
codegen = self.context.codegen()
|
||
|
|
library = codegen.create_library(self.module_name)
|
||
|
|
|
||
|
|
# Generate IR for all exported functions
|
||
|
|
flags = Flags()
|
||
|
|
flags.no_compile = True
|
||
|
|
if not self.export_python_wrap:
|
||
|
|
flags.no_cpython_wrapper = True
|
||
|
|
flags.no_cfunc_wrapper = True
|
||
|
|
if self.use_nrt:
|
||
|
|
flags.nrt = True
|
||
|
|
# Compile NRT helpers
|
||
|
|
nrt_module, _ = nrtdynmod.create_nrt_module(self.context)
|
||
|
|
library.add_ir_module(nrt_module)
|
||
|
|
|
||
|
|
for entry in self.export_entries:
|
||
|
|
cres = compile_extra(self.typing_context, self.context,
|
||
|
|
entry.function,
|
||
|
|
entry.signature.args,
|
||
|
|
entry.signature.return_type, flags,
|
||
|
|
locals={}, library=library)
|
||
|
|
|
||
|
|
func_name = cres.fndesc.llvm_func_name
|
||
|
|
llvm_func = cres.library.get_function(func_name)
|
||
|
|
|
||
|
|
if self.export_python_wrap:
|
||
|
|
llvm_func.linkage = 'internal'
|
||
|
|
wrappername = cres.fndesc.llvm_cpython_wrapper_name
|
||
|
|
wrapper = cres.library.get_function(wrappername)
|
||
|
|
wrapper.name = self._mangle_method_symbol(entry.symbol)
|
||
|
|
wrapper.linkage = 'external'
|
||
|
|
fnty = cres.target_context.call_conv.get_function_type(
|
||
|
|
cres.fndesc.restype, cres.fndesc.argtypes)
|
||
|
|
self.exported_function_types[entry] = fnty
|
||
|
|
self.function_environments[entry] = cres.environment
|
||
|
|
self.environment_gvs[entry] = cres.fndesc.env_name
|
||
|
|
else:
|
||
|
|
llvm_func.name = entry.symbol
|
||
|
|
self.dll_exports.append(entry.symbol)
|
||
|
|
|
||
|
|
if self.export_python_wrap:
|
||
|
|
wrapper_module = library.create_ir_module("wrapper")
|
||
|
|
self._emit_python_wrapper(wrapper_module)
|
||
|
|
library.add_ir_module(wrapper_module)
|
||
|
|
|
||
|
|
# Hide all functions in the DLL except those explicitly exported
|
||
|
|
library.finalize()
|
||
|
|
for fn in library.get_defined_functions():
|
||
|
|
if fn.name not in self.dll_exports:
|
||
|
|
if fn.linkage in {Linkage.private, Linkage.internal}:
|
||
|
|
# Private/Internal linkage must have "default" visibility
|
||
|
|
fn.visibility = "default"
|
||
|
|
else:
|
||
|
|
fn.visibility = 'hidden'
|
||
|
|
return library
|
||
|
|
|
||
|
|
def write_llvm_bitcode(self, output, wrap=False, **kws):
|
||
|
|
self.export_python_wrap = wrap
|
||
|
|
library = self._cull_exports()
|
||
|
|
with open(output, 'wb') as fout:
|
||
|
|
fout.write(library.emit_bitcode())
|
||
|
|
|
||
|
|
def write_native_object(self, output, wrap=False, **kws):
|
||
|
|
self.export_python_wrap = wrap
|
||
|
|
library = self._cull_exports()
|
||
|
|
with open(output, 'wb') as fout:
|
||
|
|
fout.write(library.emit_native_object())
|
||
|
|
|
||
|
|
def emit_type(self, tyobj):
|
||
|
|
ret_val = str(tyobj)
|
||
|
|
if 'int' in ret_val:
|
||
|
|
if ret_val.endswith(('8', '16', '32', '64')):
|
||
|
|
ret_val += "_t"
|
||
|
|
return ret_val
|
||
|
|
|
||
|
|
def emit_header(self, output):
|
||
|
|
fname, ext = os.path.splitext(output)
|
||
|
|
with open(fname + '.h', 'w') as fout:
|
||
|
|
fout.write(get_header())
|
||
|
|
fout.write("\n/* Prototypes */\n")
|
||
|
|
for export_entry in self.export_entries:
|
||
|
|
name = export_entry.symbol
|
||
|
|
restype = self.emit_type(export_entry.signature.return_type)
|
||
|
|
args = ", ".join(self.emit_type(argtype)
|
||
|
|
for argtype in export_entry.signature.args)
|
||
|
|
fout.write("extern %s %s(%s);\n" % (restype, name, args))
|
||
|
|
|
||
|
|
def _emit_method_array(self, llvm_module):
|
||
|
|
"""
|
||
|
|
Collect exported methods and emit a PyMethodDef array.
|
||
|
|
|
||
|
|
:returns: a pointer to the PyMethodDef array.
|
||
|
|
"""
|
||
|
|
method_defs = []
|
||
|
|
for entry in self.export_entries:
|
||
|
|
name = entry.symbol
|
||
|
|
llvm_func_name = self._mangle_method_symbol(name)
|
||
|
|
fnty = self.exported_function_types[entry]
|
||
|
|
lfunc = ir.Function(llvm_module, fnty, llvm_func_name)
|
||
|
|
|
||
|
|
method_name = self.context.insert_const_string(llvm_module, name)
|
||
|
|
method_def_const = ir.Constant.literal_struct(
|
||
|
|
(method_name,
|
||
|
|
ir.Constant.bitcast(lfunc, lt._void_star),
|
||
|
|
METH_VARARGS_AND_KEYWORDS,
|
||
|
|
NULL))
|
||
|
|
method_defs.append(method_def_const)
|
||
|
|
|
||
|
|
sentinel = ir.Constant.literal_struct([NULL, NULL, ZERO, NULL])
|
||
|
|
method_defs.append(sentinel)
|
||
|
|
method_array_init = create_constant_array(self.method_def_ty, method_defs)
|
||
|
|
method_array = cgutils.add_global_variable(llvm_module,
|
||
|
|
method_array_init.type,
|
||
|
|
'.module_methods')
|
||
|
|
method_array.initializer = method_array_init
|
||
|
|
method_array.linkage = 'internal'
|
||
|
|
method_array_ptr = ir.Constant.gep(method_array, [ZERO, ZERO])
|
||
|
|
return method_array_ptr
|
||
|
|
|
||
|
|
def _emit_environment_array(self, llvm_module, builder, pyapi):
|
||
|
|
"""
|
||
|
|
Emit an array of env_def_t structures (see modulemixin.c)
|
||
|
|
storing the pickled environment constants for each of the
|
||
|
|
exported functions.
|
||
|
|
"""
|
||
|
|
env_defs = []
|
||
|
|
for entry in self.export_entries:
|
||
|
|
env = self.function_environments[entry]
|
||
|
|
# Constants may be unhashable so avoid trying to cache them
|
||
|
|
env_def = pyapi.serialize_uncached(env.consts)
|
||
|
|
env_defs.append(env_def)
|
||
|
|
env_defs_init = create_constant_array(self.env_def_ty, env_defs)
|
||
|
|
gv = self.context.insert_unique_const(llvm_module,
|
||
|
|
'.module_environments',
|
||
|
|
env_defs_init)
|
||
|
|
return gv.gep([ZERO, ZERO])
|
||
|
|
|
||
|
|
def _emit_envgvs_array(self, llvm_module, builder, pyapi):
|
||
|
|
"""
|
||
|
|
Emit an array of Environment pointers that needs to be filled at
|
||
|
|
initialization.
|
||
|
|
"""
|
||
|
|
env_setters = []
|
||
|
|
for entry in self.export_entries:
|
||
|
|
envgv_name = self.environment_gvs[entry]
|
||
|
|
gv = self.context.declare_env_global(llvm_module, envgv_name)
|
||
|
|
envgv = gv.bitcast(lt._void_star)
|
||
|
|
env_setters.append(envgv)
|
||
|
|
|
||
|
|
env_setters_init = create_constant_array(lt._void_star, env_setters)
|
||
|
|
gv = self.context.insert_unique_const(llvm_module,
|
||
|
|
'.module_envgvs',
|
||
|
|
env_setters_init)
|
||
|
|
return gv.gep([ZERO, ZERO])
|
||
|
|
|
||
|
|
def _emit_module_init_code(self, llvm_module, builder, modobj,
|
||
|
|
method_array, env_array, envgv_array):
|
||
|
|
"""
|
||
|
|
Emit call to "external" init function, if any.
|
||
|
|
"""
|
||
|
|
if self.external_init_function:
|
||
|
|
fnty = ir.FunctionType(lt._int32,
|
||
|
|
[modobj.type, self.method_def_ptr,
|
||
|
|
self.env_def_ptr, envgv_array.type])
|
||
|
|
fn = ir.Function(llvm_module, fnty, self.external_init_function)
|
||
|
|
return builder.call(fn, [modobj, method_array, env_array,
|
||
|
|
envgv_array])
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
class ModuleCompiler(_ModuleCompiler):
|
||
|
|
|
||
|
|
_ptr_fun = lambda ret, *args: ir.PointerType(ir.FunctionType(ret, args))
|
||
|
|
|
||
|
|
#: typedef int (*visitproc)(PyObject *, void *);
|
||
|
|
visitproc_ty = _ptr_fun(lt._int8,
|
||
|
|
lt._pyobject_head_p)
|
||
|
|
|
||
|
|
#: typedef int (*inquiry)(PyObject *);
|
||
|
|
inquiry_ty = _ptr_fun(lt._int8,
|
||
|
|
lt._pyobject_head_p)
|
||
|
|
|
||
|
|
#: typedef int (*traverseproc)(PyObject *, visitproc, void *);
|
||
|
|
traverseproc_ty = _ptr_fun(lt._int8,
|
||
|
|
lt._pyobject_head_p,
|
||
|
|
visitproc_ty,
|
||
|
|
lt._void_star)
|
||
|
|
|
||
|
|
# typedef void (*freefunc)(void *)
|
||
|
|
freefunc_ty = _ptr_fun(lt._int8,
|
||
|
|
lt._void_star)
|
||
|
|
|
||
|
|
# PyObject* (*m_init)(void);
|
||
|
|
m_init_ty = _ptr_fun(lt._int8)
|
||
|
|
|
||
|
|
_char_star = lt._int8_star
|
||
|
|
|
||
|
|
#: typedef struct PyModuleDef_Base {
|
||
|
|
#: PyObject_HEAD
|
||
|
|
#: PyObject* (*m_init)(void);
|
||
|
|
#: Py_ssize_t m_index;
|
||
|
|
#: PyObject* m_copy;
|
||
|
|
#: } PyModuleDef_Base;
|
||
|
|
module_def_base_ty = ir.LiteralStructType(
|
||
|
|
(
|
||
|
|
lt._pyobject_head,
|
||
|
|
m_init_ty,
|
||
|
|
lt._llvm_py_ssize_t,
|
||
|
|
lt._pyobject_head_p
|
||
|
|
))
|
||
|
|
|
||
|
|
#: This struct holds all information that is needed to create a module object.
|
||
|
|
#: typedef struct PyModuleDef{
|
||
|
|
#: PyModuleDef_Base m_base;
|
||
|
|
#: const char* m_name;
|
||
|
|
#: const char* m_doc;
|
||
|
|
#: Py_ssize_t m_size;
|
||
|
|
#: PyMethodDef *m_methods;
|
||
|
|
#: inquiry m_reload;
|
||
|
|
#: traverseproc m_traverse;
|
||
|
|
#: inquiry m_clear;
|
||
|
|
#: freefunc m_free;
|
||
|
|
#: }PyModuleDef;
|
||
|
|
module_def_ty = ir.LiteralStructType(
|
||
|
|
(
|
||
|
|
module_def_base_ty,
|
||
|
|
_char_star,
|
||
|
|
_char_star,
|
||
|
|
lt._llvm_py_ssize_t,
|
||
|
|
_ModuleCompiler.method_def_ptr,
|
||
|
|
inquiry_ty,
|
||
|
|
traverseproc_ty,
|
||
|
|
inquiry_ty,
|
||
|
|
freefunc_ty
|
||
|
|
))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def module_create_definition(self):
|
||
|
|
"""
|
||
|
|
Return the signature and name of the Python C API function to
|
||
|
|
initialize the module.
|
||
|
|
"""
|
||
|
|
signature = ir.FunctionType(lt._pyobject_head_p,
|
||
|
|
(ir.PointerType(self.module_def_ty),
|
||
|
|
lt._int32))
|
||
|
|
|
||
|
|
name = "PyModule_Create2"
|
||
|
|
if lt._trace_refs_:
|
||
|
|
name += "TraceRefs"
|
||
|
|
|
||
|
|
return signature, name
|
||
|
|
|
||
|
|
@property
|
||
|
|
def module_init_definition(self):
|
||
|
|
"""
|
||
|
|
Return the name and signature of the module's initialization function.
|
||
|
|
"""
|
||
|
|
signature = ir.FunctionType(lt._pyobject_head_p, ())
|
||
|
|
|
||
|
|
return signature, "PyInit_" + self.module_name
|
||
|
|
|
||
|
|
def _emit_python_wrapper(self, llvm_module):
|
||
|
|
# Figure out the Python C API module creation function, and
|
||
|
|
# get a LLVM function for it.
|
||
|
|
create_module_fn = ir.Function(llvm_module, *self.module_create_definition)
|
||
|
|
create_module_fn.linkage = 'external'
|
||
|
|
|
||
|
|
# Define a constant string for the module name.
|
||
|
|
mod_name_const = self.context.insert_const_string(llvm_module,
|
||
|
|
self.module_name)
|
||
|
|
|
||
|
|
mod_def_base_init = ir.Constant.literal_struct(
|
||
|
|
(lt._pyobject_head_init, # PyObject_HEAD
|
||
|
|
ir.Constant(self.m_init_ty, None), # m_init
|
||
|
|
ir.Constant(lt._llvm_py_ssize_t, None), # m_index
|
||
|
|
ir.Constant(lt._pyobject_head_p, None), # m_copy
|
||
|
|
)
|
||
|
|
)
|
||
|
|
mod_def_base = cgutils.add_global_variable(llvm_module,
|
||
|
|
mod_def_base_init.type,
|
||
|
|
'.module_def_base')
|
||
|
|
mod_def_base.initializer = mod_def_base_init
|
||
|
|
mod_def_base.linkage = 'internal'
|
||
|
|
|
||
|
|
method_array = self._emit_method_array(llvm_module)
|
||
|
|
|
||
|
|
mod_def_init = ir.Constant.literal_struct(
|
||
|
|
(mod_def_base_init, # m_base
|
||
|
|
mod_name_const, # m_name
|
||
|
|
ir.Constant(self._char_star, None), # m_doc
|
||
|
|
ir.Constant(lt._llvm_py_ssize_t, -1), # m_size
|
||
|
|
method_array, # m_methods
|
||
|
|
ir.Constant(self.inquiry_ty, None), # m_reload
|
||
|
|
ir.Constant(self.traverseproc_ty, None), # m_traverse
|
||
|
|
ir.Constant(self.inquiry_ty, None), # m_clear
|
||
|
|
ir.Constant(self.freefunc_ty, None) # m_free
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Define a constant string for the module name.
|
||
|
|
mod_def = cgutils.add_global_variable(llvm_module, mod_def_init.type,
|
||
|
|
'.module_def')
|
||
|
|
mod_def.initializer = mod_def_init
|
||
|
|
mod_def.linkage = 'internal'
|
||
|
|
|
||
|
|
# Define the module initialization function.
|
||
|
|
mod_init_fn = ir.Function(llvm_module, *self.module_init_definition)
|
||
|
|
entry = mod_init_fn.append_basic_block('Entry')
|
||
|
|
builder = ir.IRBuilder(entry)
|
||
|
|
pyapi = self.context.get_python_api(builder)
|
||
|
|
|
||
|
|
mod = builder.call(create_module_fn,
|
||
|
|
(mod_def,
|
||
|
|
ir.Constant(lt._int32, sys.api_version)))
|
||
|
|
|
||
|
|
# Test if module has been created correctly.
|
||
|
|
# (XXX for some reason comparing with the NULL constant fails llvm
|
||
|
|
# with an assertion in pydebug mode)
|
||
|
|
with builder.if_then(cgutils.is_null(builder, mod)):
|
||
|
|
builder.ret(NULL.bitcast(mod_init_fn.type.pointee.return_type))
|
||
|
|
|
||
|
|
env_array = self._emit_environment_array(llvm_module, builder, pyapi)
|
||
|
|
envgv_array = self._emit_envgvs_array(llvm_module, builder, pyapi)
|
||
|
|
ret = self._emit_module_init_code(llvm_module, builder, mod,
|
||
|
|
method_array, env_array, envgv_array)
|
||
|
|
if ret is not None:
|
||
|
|
with builder.if_then(cgutils.is_not_null(builder, ret)):
|
||
|
|
# Init function errored out
|
||
|
|
builder.ret(ir.Constant(mod.type, None))
|
||
|
|
|
||
|
|
builder.ret(mod)
|
||
|
|
|
||
|
|
self.dll_exports.append(mod_init_fn.name)
|
||
|
|
|