"""cppjit extensions for numba"""
import re
import numba
import numba.core.cgutils as nb_cgu
import numba.core.datamodel as nb_dm
import numba.core.imputils as nb_iutils
import numba.core.registry as nb_reg
import numba.core.types as nb_types
import numba.core.typing as nb_typing
import numba.core.typing.templates as nb_tmpl
import numba.extending as nb_ext
from llvmlite import ir
import cppjit
import cppjit.reflex as cpp_refl
import cppjit.types as cpp_types
# setuptools entry point for Numba
def _init_extension():
pass
class Qualified:
default = 0
value = 1
instance = 2
ir_byte = ir.IntType(8)
ir_voidptr = ir.PointerType(ir_byte) # by convention
ir_byteptr = ir_voidptr # for clarity
ir_intptr_t = ir.IntType(cppjit.sizeof("void*") * 8)
# special case access to unboxing/boxing APIs
cppjit_as_voidptr = cppjit.addressof("Instance_AsVoidPtr")
cppjit_from_voidptr = cppjit.addressof("Instance_FromVoidPtr")
_cpp2numba = {
"void": nb_types.void,
"void*": nb_types.voidptr,
"int8_t": nb_types.int8,
"uint8_t": nb_types.uint8,
"short": nb_types.short,
"unsigned short": nb_types.ushort,
"int": nb_types.intc,
"unsigned int": nb_types.uintc,
"int32_t": nb_types.int32,
"uint32_t": nb_types.uint32,
"int64_t": nb_types.int64,
"uint64_t": nb_types.uint64,
"long": nb_types.long_,
"unsigned long": nb_types.ulong,
"long long": nb_types.longlong,
"unsigned long long": nb_types.ulonglong,
"float": nb_types.float32,
"double": nb_types.float64,
"char": nb_types.char,
"unsigned char": nb_types.uchar,
"char*": nb_types.unicode_type,
}
def resolve_std_vector(val):
return re.match(r"std::vector", val).group(1)
def resolve_const_types(val):
return re.match(r"const\s+(.+)\s*\*", val).group(1)
def cpp2numba(val):
if not isinstance(val, str):
# TODO: distinguish ptr/ref/byval
# TODO: Only metaclasses/proxies end up here since
# ref cases makes the RETURN_TYPE from reflex a string
return typeof_scope(val, nb_typing.typeof.Purpose.argument, Qualified.value)
elif val.startswith("std::vector"):
type_arr = getattr(numba, str(cpp2numba(resolve_std_vector(val))))[:]
return type_arr
elif val[-1] == "*" or val[-1] == "&":
if val.startswith("const"):
return nb_types.CPointer(cpp2numba(resolve_const_types(val)))
return nb_types.CPointer(_cpp2numba[val[:-2]])
return _cpp2numba[val]
_numba2cpp = dict()
for key, value in _cpp2numba.items():
_numba2cpp[value] = key
# prefer "int" in the case of intc over "int32_t"
_numba2cpp[nb_types.intc] = "int"
def numba2cpp(val):
if hasattr(val, "literal_type"):
val = val.literal_type
if val == nb_types.int64: # Python int
# TODO: this is only necessary until "best matching" is in place
val = nb_types.intc # more likely match candidate
elif isinstance(val, numba.types.CPointer):
return _numba2cpp[val.dtype]
elif isinstance(val, numba.types.RawPointer):
return _numba2cpp[nb_types.voidptr]
elif isinstance(val, numba.types.Array):
return "std::vector"
elif isinstance(val, CppClassNumbaType):
return val._scope.__cpp_name__
else:
try:
return _numba2cpp[val]
except:
raise RuntimeError("Type mapping failed from Numba to C++ for ", val)
def numba_arg_convertor(args):
args_cpp = []
for arg in list(args):
# If the user explicitly passes an argument using numba CPointer, the regex match is used
# to detect the pass by reference since the dispatcher always returns typeref[val*]
match = re.search(r"typeref\[(.*?)\*\]", str(arg))
if match:
literal_val = match.group(1)
arg_type = numba.typeof(eval(literal_val))
args_cpp.append(to_ref(numba2cpp(arg_type)))
else:
args_cpp.append(numba2cpp(arg))
return tuple(args_cpp)
def to_ref(type_list):
ref_list = []
for l in type_list: # noqa: E741
ref_list.append(l + " &")
return ref_list
# TODO: looks like Numba treats unsigned types as signed when lowering,
# which seems to work as they're just reinterpret_casts
_cpp2ir = {
"char *": ir_byteptr,
"int8_t": ir.IntType(8),
"uint8_t": ir.IntType(8),
"short": ir.IntType(nb_types.short.bitwidth),
"unsigned short": ir.IntType(nb_types.ushort.bitwidth),
"int": ir.IntType(nb_types.intc.bitwidth),
"unsigned int": ir.IntType(nb_types.uintc.bitwidth),
"int32_t": ir.IntType(32),
"uint32_t": ir.IntType(32),
"int64_t": ir.IntType(64),
"uint64_t": ir.IntType(64),
"long": ir.IntType(nb_types.long_.bitwidth),
"unsigned long": ir.IntType(nb_types.ulong.bitwidth),
"long long": ir.IntType(nb_types.longlong.bitwidth),
"unsigned long long": ir.IntType(nb_types.ulonglong.bitwidth),
"float": ir.FloatType(),
"double": ir.DoubleType(),
}
def cpp2ir(val):
try:
return _cpp2ir[val]
except KeyError:
if val.startswith("std::vector"):
## TODO should be possible to obtain the vector length from the CPPDataMember val
type_arr = ir.VectorType(cpp2ir(resolve_std_vector(val)), 3)
return type_arr
elif val != "char *" and val[-1] == "*":
if val.startswith("const"):
return ir.PointerType(cpp2ir(resolve_const_types(val)))
type_2 = _cpp2ir[val[:-2]]
return ir.PointerType(type_2)
#
# C++ function pointer -> Numba
#
class CppFunctionNumbaType(nb_types.Callable):
targetdescr = nb_reg.cpu_target
requires_gil = False
def __init__(self, func, is_method=False):
super(CppFunctionNumbaType, self).__init__("CppFunction(%s)" % str(func))
self.sig = None
self._func = func
self._is_method = is_method
self._signatures = list()
self._impl_keys = dict()
self._arg_set_matched = tuple()
self.ret_type = None
def is_precise(self):
return True # by definition
def get_call_type(self, context, args, kwds):
try:
return self._impl_keys[args].sig
except KeyError:
pass
ol = CppFunctionNumbaType(
self._func.__overload__(numba_arg_convertor(args)), self._is_method
)
thistype = None
if self._is_method:
thistype = nb_types.voidptr
self.ret_type = cpp2numba(ol._func.__cpp_reflex__(cpp_refl.RETURN_TYPE))
ol.sig = nb_typing.Signature(
return_type=self.ret_type, args=args, recvr=thistype
)
extsig = ol.sig
if self._is_method:
self.ret_type = ol.sig.return_type
args = (nb_types.voidptr, *args)
extsig = nb_typing.Signature(
return_type=ol.sig.return_type, args=args, recvr=None
)
self._impl_keys[args] = ol
self._arg_set_matched = numba_arg_convertor(args)
@nb_iutils.lower_builtin(ol, *args)
def lower_external_call(
context,
builder,
sig,
args,
ty=nb_types.ExternalFunctionPointer(extsig, ol.get_pointer),
pyval=self._func,
is_method=self._is_method,
):
ptrty = context.get_function_pointer_type(ty)
ptrval = context.add_dynamic_addr(
builder, ty.get_pointer(pyval), info=str(pyval)
)
fptr = builder.bitcast(ptrval, ptrty)
return context.call_function_pointer(builder, fptr, args)
return ol.sig
def get_call_signatures(self):
return list(self._signatures), False
def get_impl_key(self, sig):
return self._impl_keys[sig.args]
# TODO: Remove the redundancy of __overload__ matching and use this function
# to only obtain the address given the matched overload
def get_pointer(self, func):
if func is None:
func = self._func
ol = func.__overload__(numba_arg_convertor(self.sig.args))
address = cppjit.addressof(ol)
if not address:
raise RuntimeError("unresolved address for %s" % str(ol))
return address
@property
def key(self):
return self._func
@nb_ext.typeof_impl.register(cpp_types.Function)
def typeof_function(val, c):
return CppFunctionNumbaType(val)
@nb_ext.typeof_impl.register(cpp_types.Template)
def typeof_template(val, c):
if hasattr(val, "__overload__"):
return CppFunctionNumbaType(val)
raise RuntimeError("only function templates supported")
@nb_ext.register_model(CppFunctionNumbaType)
class CppFunctionModel(nb_dm.models.PrimitiveModel):
def __init__(self, dmm, fe_type):
# the function pointer of this overload can not be exactly typed, but
# only the storage size is relevant, so simply use a void*
be_type = ir.PointerType(dmm.lookup(nb_types.void).get_value_type())
super(CppFunctionModel, self).__init__(dmm, fe_type, be_type)
@nb_iutils.lower_constant(CppFunctionNumbaType)
def constant_function_pointer(context, builder, ty, pyval):
# TODO: needs to exist for the proper flow, but why? The lowering of the
# actual overload is handled dynamically.
return
#
# C++ method / data member -> Numba
#
class CppDataMemberInfo(object):
__slots__ = ["f_name", "f_offset", "f_nbtype", "f_irtype"]
def __init__(self, name, offset, cpptype):
self.f_name = name
self.f_offset = offset
self.f_nbtype = cpp2numba(cpptype)
self.f_irtype = cpp2ir(cpptype)
#
# C++ class -> Numba
#
class CppClassNumbaType(CppFunctionNumbaType):
def __init__(self, scope, qualifier):
addr = None
cppinstance_val = None
if qualifier == Qualified.instance:
addr = cppjit.addressof(scope)
cppinstance_val = scope
scope = type(scope)
qualifier = Qualified.default
super(CppClassNumbaType, self).__init__(scope.__init__)
self.name = "CppClass(%s)" % scope.__cpp_name__ # overrides value in Type
self._scope = scope
self._qualifier = qualifier
self._cppinstanceval = cppinstance_val
self._addr = addr
def get_scope(self):
return self._scope
def get_qualifier(self):
return self._qualifier
def get_call_type(self, context, args, kwds):
sig = super(CppClassNumbaType, self).get_call_type(context, args, kwds)
self.sig = sig
return sig
def is_precise(self):
return True
@property
def key(self):
return (self._scope, self._qualifier)
@nb_tmpl.infer_getattr
class CppClassFieldResolver(nb_tmpl.AttributeTemplate):
key = CppClassNumbaType
def generic_resolve(self, typ, attr):
ft = typ.__dict__.get(attr, None)
if ft is not None:
return ft
try:
f = getattr(typ._scope, attr)
if isinstance(f, cpp_types.Function):
ft = CppFunctionNumbaType(f, is_method=True)
except AttributeError:
pass
try:
f = typ._scope.__dict__[attr]
if isinstance(f, cpp_types.DataMember):
ct = f.__cpp_reflex__(cpp_refl.TYPE)
ft = cpp2numba(ct)
except AttributeError:
pass
if ft is not None:
typ.__dict__[attr] = ft
return ft
@nb_iutils.lower_getattr_generic(CppClassNumbaType)
def cppclass_getattr_impl(context, builder, typ, val, attr):
# TODO: the following relies on the fact that numba will first lower the
# field access, then immediately lower the call; and that the `val` loads
# the struct representing the C++ object. Neither need be stable.
if attr in typ._scope.__dict__ and isinstance(
typ._scope.__dict__[attr], cpp_types.DataMember
):
dm = typ._scope.__dict__[attr]
ct = dm.__cpp_reflex__(cpp_refl.TYPE)
offset = dm.__cpp_reflex__(cpp_refl.OFFSET)
q = typ.get_qualifier()
if q == Qualified.default:
llval = builder.bitcast(val, ir_byteptr)
pfc = builder.gep(llval, [ir.Constant(ir_intptr_t, offset)])
pf = builder.bitcast(pfc, ir.PointerType(cpp2ir(ct)))
return builder.load(pf)
elif q == Qualified.value:
model = nb_dm.default_manager.lookup(typ)
return model.get(builder, val, attr)
else:
assert not "unknown qualified type"
# TODO: easier with inttoptr and ptrtoint (cgutils.pointer_add)?
llval = builder.bitcast(val, ir_byteptr)
pfc = builder.gep(llval, [ir.Constant(ir_intptr_t, offset)])
pf = builder.bitcast(pfc, ir.PointerType(cpp2ir(ct)))
return builder.load(pf)
# assume this is a method
q = typ.get_qualifier()
if q == Qualified.default:
return builder.bitcast(val, ir_voidptr)
elif q == Qualified.value:
return None
assert not "unknown qualified type"
return None
class ImplAggregateValueModel(nb_dm.models.StructModel):
def get(self, builder, val, pos):
"""Get a field at the given position/field name"""
if isinstance(pos, str):
pos = self.get_field_position(pos)
# Use the offsets for direct addressing, rather than getting the elements
# from the struct type.
dmi = self._data_members[pos]
stack = nb_cgu.alloca_once(builder, self.get_data_type())
builder.store(val, stack)
llval = builder.bitcast(stack, ir_byteptr)
pfc = builder.gep(llval, [ir.Constant(ir_intptr_t, dmi.f_offset)])
pf = builder.bitcast(pfc, ir.PointerType(dmi.f_irtype))
return builder.load(pf)
class ImplClassValueModel(ImplAggregateValueModel):
# TODO : Should the address have to be passed here and stored in meminfo
# value: representation inside function body. Maybe stored in stack.
# The representation here are flexible.
def get_value_type(self):
return self.get_data_type()
# data: representation used when storing into containers (e.g. arrays).
def get_data_type(self):
# The struct model relies on data being a POD, but for C++ objects, there
# can be hidden data (e.g. vtable, thunks, or simply private members), and
# the alignment of Cling and Numba also need not be the same. Therefore, the
# struct is split in a series of byte members to get the total size right
# and to allow addressing at the correct offsets.
if self._data_type is None:
self._data_type = ir.LiteralStructType(
[ir_byte for i in range(self._sizeof)], packed=True
)
return self._data_type
# return: representation used for return argument.
def get_return_type(self):
return self.get_data_type()
scope_numbatypes = (dict(), dict())
@nb_ext.typeof_impl.register(cpp_types.Scope)
def typeof_scope(val, c, q=Qualified.default):
is_instance = False
cppinstance_val = None
if q == Qualified.instance:
cppinstance_val = val
val = type(val)
q = Qualified.default
is_instance = True
global scope_numbatypes
try:
return scope_numbatypes[q][val]
except KeyError:
pass
if val.__cpp_reflex__(cpp_refl.IS_NAMESPACE):
cnt = nb_types.Module(val)
scope_numbatypes[Qualified.default][val] = cnt
return cnt
class ImplClassType(CppClassNumbaType):
pass
if is_instance:
cnt = ImplClassType(cppinstance_val, Qualified.instance)
else:
cnt = ImplClassType(val, q)
scope_numbatypes[q][val] = cnt
# declare data members to Numba
data_members = list()
member_methods = dict()
for name, field in val.__dict__.items():
if isinstance(field, cpp_types.DataMember):
data_members.append(
CppDataMemberInfo(
name,
field.__cpp_reflex__(cpp_refl.OFFSET),
field.__cpp_reflex__(cpp_refl.TYPE),
)
)
elif isinstance(field, cpp_types.Function):
member_methods[name] = field.__cpp_reflex__(cpp_refl.RETURN_TYPE)
# TODO: this refresh is needed b/c the scope type is registered as a
# callable after the tracing started; no idea of the side-effects ...
nb_reg.cpu_target.typing_context.refresh()
# create a model description for Numba
if q == Qualified.default:
@nb_ext.register_model(ImplClassType)
class ImplClassModel(nb_dm.models.StructModel):
def __init__(self, dmm, fe_type):
self._data_members = data_members
self._member_methods = member_methods
# TODO: eventually we need not derive from StructModel
members = [(dmi.f_name, dmi.f_nbtype) for dmi in data_members]
nb_dm.models.StructModel.__init__(self, dmm, fe_type, members)
# proxies are always accessed by pointer, which are not composites
def traverse(self, builder):
return []
def traverse_models(self):
return []
def traverse_types(self):
return [self._fe_type] # from StructModel
# data: representation used when storing into containers (e.g. arrays).
# TODO ...
# value: representation inside function body. Maybe stored in stack.
# The representation here are flexible.
def get_value_type(self):
# the C++ object, b/c through a proxy, is always accessed by pointer; it is
# represented as a pointer to POD to allow indexing by Numba for data member
# type checking, but the address offsetting for loading data member values is
# independent (see get(), below), so the exact layout need not match a POD
# TODO: this doesn't work for real PODs, b/c those are unpacked into their elements
# and passed through registers
return ir.PointerType(super(ImplClassModel, self).get_value_type())
# argument: representation used for function argument. Needs to be builtin type,
# but unlike other Numba composites, C++ proxies are not flattened.
def get_argument_type(self):
return self.get_value_type()
def as_argument(self, builder, value):
return value
def from_argument(self, builder, value):
return value
# return: representation used for return argument.
# TODO ...
# access to public data members
def get(self, builder, val, pos):
"""Get a field at the given position/field name"""
if isinstance(pos, str):
pos = self.get_field_position(pos)
dmi = self._data_members[pos]
llval = builder.bitcast(val, ir_byteptr)
pfc = builder.gep(llval, [ir.Constant(ir_intptr_t, dmi.f_offset)])
pf = builder.bitcast(pfc, ir.PointerType(dmi.f_irtype))
return builder.load(pf)
elif q == Qualified.value:
if val.__cpp_reflex__(cpp_refl.IS_AGGREGATE):
@nb_ext.register_model(ImplClassType)
class ImplClassModel(ImplAggregateValueModel):
pass
else:
@nb_ext.register_model(ImplClassType)
class ImplClassModel(ImplClassValueModel):
pass
def init(self, dmm, fe_type, sz=cppjit.sizeof(val)):
self._data_members = data_members
self._member_methods = member_methods
self._sizeof = sz
# TODO: this code exists purely to be able to use the indexing and hierarchy
# of the base class StructModel, which isn't much of a reason
members = [(dmi.f_name, dmi.f_nbtype) for dmi in data_members]
nb_dm.models.StructModel.__init__(self, dmm, fe_type, members)
ImplClassModel.__init__ = init
else:
assert not "unknown qualified type"
# Python proxy unwrapping for arguments into the Numba trace
@nb_ext.unbox(ImplClassType)
def unbox_instance(typ, obj, c):
global cppjit_as_voidptr
ptrty = ir.PointerType(ir.FunctionType(ir_voidptr, [ir_voidptr]))
ptrval = c.context.add_dynamic_addr(
c.builder, cppjit_as_voidptr, info="Instance_AsVoidPtr"
)
fp = c.builder.bitcast(ptrval, ptrty)
vptr = c.context.call_function_pointer(c.builder, fp, [obj])
model = nb_dm.default_manager.lookup(typ)
pobj = c.builder.bitcast(vptr, model.get_argument_type())
return nb_ext.NativeValue(pobj, is_error=None, cleanup=None)
def make_implclass(context, builder, typ, **kwargs):
return nb_cgu.create_struct_proxy(typ)(context, builder, **kwargs)
# C++ object to Python proxy wrapping for returns from Numba trace
@nb_ext.box(ImplClassType)
def box_instance(typ, val, c):
assert not "requires object model and passing of intact object, not memberwise copy"
global cppjit_from_voidptr
if isinstance(val, ir.Constant):
if val.constant == ir.Undefined:
assert not "Value passed to instance boxing is undefined"
return NULL # noqa: F821
implclass = make_implclass(c.context, c.builder, typ)
classobj = c.pyapi.unserialize(c.pyapi.serialize_object(cpp_types.Instance))
box_list = []
model = implclass._datamodel # noqa: F841
cfr = CppClassFieldResolver(c.context)
for i in typ._scope.__dict__:
if isinstance(cfr.generic_resolve(typ, i), nb_types.Type):
box_list.append(
c.box(cfr.generic_resolve(typ, i), getattr(implclass, i))
)
box_res = c.pyapi.call_function_objargs(classobj, tuple(box_list))
# Required for nopython mode, numba nrt requres each member box call to decref
# since it steals the reference
for i in box_list:
c.pyapi.decref(i)
return box_res
return cnt
#
# C++ instance -> Numba
#
@nb_ext.typeof_impl.register(cpp_types.Instance)
def typeof_instance(val, c):
global scope_numbatypes
try:
return scope_numbatypes[Qualified.default][type(val)]
except KeyError:
pass
# Pass the val itself to obtain Cling address of the CPPInstance for reference to C++ objects
return typeof_scope(val, c, Qualified.instance)