Skip to content
Permalink
master
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
from functools import partial
from collections import deque
from llvmlite import ir
from numba.core.datamodel.registry import register_default
from numba.core import types, cgutils
from numba.np import numpy_support
class DataModel(object):
"""
DataModel describe how a FE type is represented in the LLVM IR at
different contexts.
Contexts are:
- value: representation inside function body. Maybe stored in stack.
The representation here are flexible.
- data: representation used when storing into containers (e.g. arrays).
- argument: representation used for function argument. All composite
types are unflattened into multiple primitive types.
- return: representation used for return argument.
Throughput the compiler pipeline, a LLVM value is usually passed around
in the "value" representation. All "as_" prefix function converts from
"value" representation. All "from_" prefix function converts to the
"value" representation.
"""
def __init__(self, dmm, fe_type):
self._dmm = dmm
self._fe_type = fe_type
@property
def fe_type(self):
return self._fe_type
def get_value_type(self):
raise NotImplementedError(self)
def get_data_type(self):
return self.get_value_type()
def get_argument_type(self):
"""Return a LLVM type or nested tuple of LLVM type
"""
return self.get_value_type()
def get_return_type(self):
return self.get_value_type()
def as_data(self, builder, value):
raise NotImplementedError(self)
def as_argument(self, builder, value):
"""
Takes one LLVM value
Return a LLVM value or nested tuple of LLVM value
"""
raise NotImplementedError(self)
def as_return(self, builder, value):
raise NotImplementedError(self)
def from_data(self, builder, value):
raise NotImplementedError(self)
def from_argument(self, builder, value):
"""
Takes a LLVM value or nested tuple of LLVM value
Returns one LLVM value
"""
raise NotImplementedError(self)
def from_return(self, builder, value):
raise NotImplementedError(self)
def load_from_data_pointer(self, builder, ptr, align=None):
"""
Load value from a pointer to data.
This is the default implementation, sufficient for most purposes.
"""
return self.from_data(builder, builder.load(ptr, align=align))
def traverse(self, builder):
"""
Traverse contained members.
Returns a iterable of contained (types, getters).
Each getter is a one-argument function accepting a LLVM value.
"""
return []
def traverse_models(self):
"""
Recursively list all models involved in this model.
"""
return [self._dmm[t] for t in self.traverse_types()]
def traverse_types(self):
"""
Recursively list all frontend types involved in this model.
"""
types = [self._fe_type]
queue = deque([self])
while len(queue) > 0:
dm = queue.popleft()
for i_dm in dm.inner_models():
if i_dm._fe_type not in types:
queue.append(i_dm)
types.append(i_dm._fe_type)
return types
def inner_models(self):
"""
List all *inner* models.
"""
return []
def get_nrt_meminfo(self, builder, value):
"""
Returns the MemInfo object or None if it is not tracked.
It is only defined for types.meminfo_pointer
"""
return None
def has_nrt_meminfo(self):
return False
def contains_nrt_meminfo(self):
"""
Recursively check all contained types for need for NRT meminfo.
"""
return any(model.has_nrt_meminfo() for model in self.traverse_models())
def _compared_fields(self):
return (type(self), self._fe_type)
def __hash__(self):
return hash(tuple(self._compared_fields()))
def __eq__(self, other):
if type(self) is type(other):
return self._compared_fields() == other._compared_fields()
else:
return False
def __ne__(self, other):
return not self.__eq__(other)
@register_default(types.Omitted)
class OmittedArgDataModel(DataModel):
"""
A data model for omitted arguments. Only the "argument" representation
is defined, other representations raise a NotImplementedError.
"""
# Omitted arguments don't produce any LLVM function argument.
def get_argument_type(self):
return ()
def as_argument(self, builder, val):
return ()
def from_argument(self, builder, val):
assert val == (), val
return None
@register_default(types.Boolean)
@register_default(types.BooleanLiteral)
class BooleanModel(DataModel):
_bit_type = ir.IntType(1)
_byte_type = ir.IntType(8)
def get_value_type(self):
return self._bit_type
def get_data_type(self):
return self._byte_type
def get_return_type(self):
return self.get_data_type()
def get_argument_type(self):
return self.get_data_type()
def as_data(self, builder, value):
return builder.zext(value, self.get_data_type())
def as_argument(self, builder, value):
return self.as_data(builder, value)
def as_return(self, builder, value):
return self.as_data(builder, value)
def from_data(self, builder, value):
ty = self.get_value_type()
resalloca = cgutils.alloca_once(builder, ty)
cond = builder.icmp_unsigned('==', value, value.type(0))
with builder.if_else(cond) as (then, otherwise):
with then:
builder.store(ty(0), resalloca)
with otherwise:
builder.store(ty(1), resalloca)
return builder.load(resalloca)
def from_argument(self, builder, value):
return self.from_data(builder, value)
def from_return(self, builder, value):
return self.from_data(builder, value)
class PrimitiveModel(DataModel):
"""A primitive type can be represented natively in the target in all
usage contexts.
"""
def __init__(self, dmm, fe_type, be_type):
super(PrimitiveModel, self).__init__(dmm, fe_type)
self.be_type = be_type
def get_value_type(self):
return self.be_type
def as_data(self, builder, value):
return value
def as_argument(self, builder, value):
return value
def as_return(self, builder, value):
return value
def from_data(self, builder, value):
return value
def from_argument(self, builder, value):
return value
def from_return(self, builder, value):
return value
class ProxyModel(DataModel):
"""
Helper class for models which delegate to another model.
"""
def get_value_type(self):
return self._proxied_model.get_value_type()
def get_data_type(self):
return self._proxied_model.get_data_type()
def get_return_type(self):
return self._proxied_model.get_return_type()
def get_argument_type(self):
return self._proxied_model.get_argument_type()
def as_data(self, builder, value):
return self._proxied_model.as_data(builder, value)
def as_argument(self, builder, value):
return self._proxied_model.as_argument(builder, value)
def as_return(self, builder, value):
return self._proxied_model.as_return(builder, value)
def from_data(self, builder, value):
return self._proxied_model.from_data(builder, value)
def from_argument(self, builder, value):
return self._proxied_model.from_argument(builder, value)
def from_return(self, builder, value):
return self._proxied_model.from_return(builder, value)
@register_default(types.EnumMember)
@register_default(types.IntEnumMember)
class EnumModel(ProxyModel):
"""
Enum members are represented exactly like their values.
"""
def __init__(self, dmm, fe_type):
super(EnumModel, self).__init__(dmm, fe_type)
self._proxied_model = dmm.lookup(fe_type.dtype)
@register_default(types.Opaque)
@register_default(types.PyObject)
@register_default(types.RawPointer)
@register_default(types.NoneType)
@register_default(types.StringLiteral)
@register_default(types.EllipsisType)
@register_default(types.Function)
@register_default(types.Type)
@register_default(types.Object)
@register_default(types.Module)
@register_default(types.Phantom)
@register_default(types.ContextManager)
@register_default(types.Dispatcher)
@register_default(types.ObjModeDispatcher)
@register_default(types.ExceptionClass)
@register_default(types.Dummy)
@register_default(types.ExceptionInstance)
@register_default(types.ExternalFunction)
@register_default(types.EnumClass)
@register_default(types.IntEnumClass)
@register_default(types.NumberClass)
@register_default(types.TypeRef)
@register_default(types.NamedTupleClass)
@register_default(types.DType)
@register_default(types.RecursiveCall)
@register_default(types.MakeFunctionLiteral)
@register_default(types.Poison)
class OpaqueModel(PrimitiveModel):
"""
Passed as opaque pointers
"""
_ptr_type = ir.IntType(8).as_pointer()
def __init__(self, dmm, fe_type):
be_type = self._ptr_type
super(OpaqueModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.MemInfoPointer)
class MemInfoModel(OpaqueModel):
def inner_models(self):
return [self._dmm.lookup(self._fe_type.dtype)]
def has_nrt_meminfo(self):
return True
def get_nrt_meminfo(self, builder, value):
return value
@register_default(types.Integer)
@register_default(types.IntegerLiteral)
class IntegerModel(PrimitiveModel):
def __init__(self, dmm, fe_type):
be_type = ir.IntType(fe_type.bitwidth)
super(IntegerModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.Float)
class FloatModel(PrimitiveModel):
def __init__(self, dmm, fe_type):
if fe_type == types.float32:
be_type = ir.FloatType()
elif fe_type == types.float64:
be_type = ir.DoubleType()
else:
raise NotImplementedError(fe_type)
super(FloatModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.CPointer)
class PointerModel(PrimitiveModel):
def __init__(self, dmm, fe_type):
self._pointee_model = dmm.lookup(fe_type.dtype)
self._pointee_be_type = self._pointee_model.get_data_type()
be_type = self._pointee_be_type.as_pointer()
super(PointerModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.EphemeralPointer)
class EphemeralPointerModel(PointerModel):
def get_data_type(self):
return self._pointee_be_type
def as_data(self, builder, value):
value = builder.load(value)
return self._pointee_model.as_data(builder, value)
def from_data(self, builder, value):
raise NotImplementedError("use load_from_data_pointer() instead")
def load_from_data_pointer(self, builder, ptr, align=None):
return builder.bitcast(ptr, self.get_value_type())
@register_default(types.EphemeralArray)
class EphemeralArrayModel(PointerModel):
def __init__(self, dmm, fe_type):
super(EphemeralArrayModel, self).__init__(dmm, fe_type)
self._data_type = ir.ArrayType(self._pointee_be_type,
self._fe_type.count)
def get_data_type(self):
return self._data_type
def as_data(self, builder, value):
values = [builder.load(cgutils.gep_inbounds(builder, value, i))
for i in range(self._fe_type.count)]
return cgutils.pack_array(builder, values)
def from_data(self, builder, value):
raise NotImplementedError("use load_from_data_pointer() instead")
def load_from_data_pointer(self, builder, ptr, align=None):
return builder.bitcast(ptr, self.get_value_type())
@register_default(types.ExternalFunctionPointer)
class ExternalFuncPointerModel(PrimitiveModel):
def __init__(self, dmm, fe_type):
sig = fe_type.sig
# Since the function is non-Numba, there is no adaptation
# of arguments and return value, hence get_value_type().
retty = dmm.lookup(sig.return_type).get_value_type()
args = [dmm.lookup(t).get_value_type() for t in sig.args]
be_type = ir.PointerType(ir.FunctionType(retty, args))
super(ExternalFuncPointerModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.UniTuple)
@register_default(types.NamedUniTuple)
@register_default(types.StarArgUniTuple)
class UniTupleModel(DataModel):
def __init__(self, dmm, fe_type):
super(UniTupleModel, self).__init__(dmm, fe_type)
self._elem_model = dmm.lookup(fe_type.dtype)
self._count = len(fe_type)
self._value_type = ir.ArrayType(self._elem_model.get_value_type(),
self._count)
self._data_type = ir.ArrayType(self._elem_model.get_data_type(),
self._count)
def get_value_type(self):
return self._value_type
def get_data_type(self):
return self._data_type
def get_return_type(self):
return self.get_value_type()
def get_argument_type(self):
return (self._elem_model.get_argument_type(),) * self._count
def as_argument(self, builder, value):
out = []
for i in range(self._count):
v = builder.extract_value(value, [i])
v = self._elem_model.as_argument(builder, v)
out.append(v)
return out
def from_argument(self, builder, value):
out = ir.Constant(self.get_value_type(), ir.Undefined)
for i, v in enumerate(value):
v = self._elem_model.from_argument(builder, v)
out = builder.insert_value(out, v, [i])
return out
def as_data(self, builder, value):
out = ir.Constant(self.get_data_type(), ir.Undefined)
for i in range(self._count):
val = builder.extract_value(value, [i])
dval = self._elem_model.as_data(builder, val)
out = builder.insert_value(out, dval, [i])
return out
def from_data(self, builder, value):
out = ir.Constant(self.get_value_type(), ir.Undefined)
for i in range(self._count):
val = builder.extract_value(value, [i])
dval = self._elem_model.from_data(builder, val)
out = builder.insert_value(out, dval, [i])
return out
def as_return(self, builder, value):
return value
def from_return(self, builder, value):
return value
def traverse(self, builder):
def getter(i, value):
return builder.extract_value(value, i)
return [(self._fe_type.dtype, partial(getter, i))
for i in range(self._count)]
def inner_models(self):
return [self._elem_model]
class CompositeModel(DataModel):
"""Any model that is composed of multiple other models should subclass from
this.
"""
pass
class StructModel(CompositeModel):
_value_type = None
_data_type = None
def __init__(self, dmm, fe_type, members):
super(StructModel, self).__init__(dmm, fe_type)
if members:
self._fields, self._members = zip(*members)
else:
self._fields = self._members = ()
self._models = tuple([self._dmm.lookup(t) for t in self._members])
def get_member_fe_type(self, name):
"""
StructModel-specific: get the Numba type of the field named *name*.
"""
pos = self.get_field_position(name)
return self._members[pos]
def get_value_type(self):
if self._value_type is None:
self._value_type = ir.LiteralStructType([t.get_value_type()
for t in self._models])
return self._value_type
def get_data_type(self):
if self._data_type is None:
self._data_type = ir.LiteralStructType([t.get_data_type()
for t in self._models])
return self._data_type
def get_argument_type(self):
return tuple([t.get_argument_type() for t in self._models])
def get_return_type(self):
return self.get_data_type()
def _as(self, methname, builder, value):
extracted = []
for i, dm in enumerate(self._models):
extracted.append(getattr(dm, methname)(builder,
self.get(builder, value, i)))
return tuple(extracted)
def _from(self, methname, builder, value):
struct = ir.Constant(self.get_value_type(), ir.Undefined)
for i, (dm, val) in enumerate(zip(self._models, value)):
v = getattr(dm, methname)(builder, val)
struct = self.set(builder, struct, v, i)
return struct
def as_data(self, builder, value):
"""
Converts the LLVM struct in `value` into a representation suited for
storing into arrays.
Note
----
Current implementation rarely changes how types are represented for
"value" and "data". This is usually a pointless rebuild of the
immutable LLVM struct value. Luckily, LLVM optimization removes all
redundancy.
Sample usecase: Structures nested with pointers to other structures
that can be serialized into a flat representation when storing into
array.
"""
elems = self._as("as_data", builder, value)
struct = ir.Constant(self.get_data_type(), ir.Undefined)
for i, el in enumerate(elems):
struct = builder.insert_value(struct, el, [i])
return struct
def from_data(self, builder, value):
"""
Convert from "data" representation back into "value" representation.
Usually invoked when loading from array.
See notes in `as_data()`
"""
vals = [builder.extract_value(value, [i])
for i in range(len(self._members))]
return self._from("from_data", builder, vals)
def load_from_data_pointer(self, builder, ptr, align=None):
values = []
for i, model in enumerate(self._models):
elem_ptr = cgutils.gep_inbounds(builder, ptr, 0, i)
val = model.load_from_data_pointer(builder, elem_ptr, align)
values.append(val)
struct = ir.Constant(self.get_value_type(), ir.Undefined)
for i, val in enumerate(values):
struct = self.set(builder, struct, val, i)
return struct
def as_argument(self, builder, value):
return self._as("as_argument", builder, value)
def from_argument(self, builder, value):
return self._from("from_argument", builder, value)
def as_return(self, builder, value):
elems = self._as("as_data", builder, value)
struct = ir.Constant(self.get_data_type(), ir.Undefined)
for i, el in enumerate(elems):
struct = builder.insert_value(struct, el, [i])
return struct
def from_return(self, builder, value):
vals = [builder.extract_value(value, [i])
for i in range(len(self._members))]
return self._from("from_data", builder, vals)
def get(self, builder, val, pos):
"""Get a field at the given position or the fieldname
Args
----
builder:
LLVM IRBuilder
val:
value to be inserted
pos: int or str
field index or field name
Returns
-------
Extracted value
"""
if isinstance(pos, str):
pos = self.get_field_position(pos)
return builder.extract_value(val, [pos],
name="extracted." + self._fields[pos])
def set(self, builder, stval, val, pos):
"""Set a field at the given position or the fieldname
Args
----
builder:
LLVM IRBuilder
stval:
LLVM struct value
val:
value to be inserted
pos: int or str
field index or field name
Returns
-------
A new LLVM struct with the value inserted
"""
if isinstance(pos, str):
pos = self.get_field_position(pos)
return builder.insert_value(stval, val, [pos],
name="inserted." + self._fields[pos])
def get_field_position(self, field):
try:
return self._fields.index(field)
except ValueError:
raise KeyError("%s does not have a field named %r"
% (self.__class__.__name__, field))
@property
def field_count(self):
return len(self._fields)
def get_type(self, pos):
"""Get the frontend type (numba type) of a field given the position
or the fieldname
Args
----
pos: int or str
field index or field name
"""
if isinstance(pos, str):
pos = self.get_field_position(pos)
return self._members[pos]
def get_model(self, pos):
"""
Get the datamodel of a field given the position or the fieldname.
Args
----
pos: int or str
field index or field name
"""
return self._models[pos]
def traverse(self, builder):
def getter(k, value):
if value.type != self.get_value_type():
args = self.get_value_type(), value.type
raise TypeError("expecting {0} but got {1}".format(*args))
return self.get(builder, value, k)
return [(self.get_type(k), partial(getter, k)) for k in self._fields]
def inner_models(self):
return self._models
@register_default(types.Complex)
class ComplexModel(StructModel):
_element_type = NotImplemented
def __init__(self, dmm, fe_type):
members = [
('real', fe_type.underlying_float),
('imag', fe_type.underlying_float),
]
super(ComplexModel, self).__init__(dmm, fe_type, members)
@register_default(types.LiteralList)
@register_default(types.LiteralStrKeyDict)
@register_default(types.Tuple)
@register_default(types.NamedTuple)
@register_default(types.StarArgTuple)
class TupleModel(StructModel):
def __init__(self, dmm, fe_type):
members = [('f' + str(i), t) for i, t in enumerate(fe_type)]
super(TupleModel, self).__init__(dmm, fe_type, members)
@register_default(types.UnionType)
class UnionModel(StructModel):
def __init__(self, dmm, fe_type):
members = [
('tag', types.uintp),
# XXX: it should really be a MemInfoPointer(types.voidptr)
('payload', types.Tuple.from_types(fe_type.types)),
]
super(UnionModel, self).__init__(dmm, fe_type, members)
@register_default(types.Pair)
class PairModel(StructModel):
def __init__(self, dmm, fe_type):
members = [('first', fe_type.first_type),
('second', fe_type.second_type)]
super(PairModel, self).__init__(dmm, fe_type, members)
@register_default(types.ListPayload)
class ListPayloadModel(StructModel):
def __init__(self, dmm, fe_type):
# The fields are mutable but the payload is always manipulated
# by reference. This scheme allows mutations of an array to
# be seen by its iterators.
members = [
('size', types.intp),
('allocated', types.intp),
# This member is only used only for reflected lists
('dirty', types.boolean),
# Actually an inlined var-sized array
('data', fe_type.container.dtype),
]
super(ListPayloadModel, self).__init__(dmm, fe_type, members)
@register_default(types.List)
class ListModel(StructModel):
def __init__(self, dmm, fe_type):
payload_type = types.ListPayload(fe_type)
members = [
# The meminfo data points to a ListPayload
('meminfo', types.MemInfoPointer(payload_type)),
# This member is only used only for reflected lists
('parent', types.pyobject),
]
super(ListModel, self).__init__(dmm, fe_type, members)
@register_default(types.ListIter)
class ListIterModel(StructModel):
def __init__(self, dmm, fe_type):
payload_type = types.ListPayload(fe_type.container)
members = [
# The meminfo data points to a ListPayload (shared with the
# original list object)
('meminfo', types.MemInfoPointer(payload_type)),
('index', types.EphemeralPointer(types.intp)),
]
super(ListIterModel, self).__init__(dmm, fe_type, members)
@register_default(types.SetEntry)
class SetEntryModel(StructModel):
def __init__(self, dmm, fe_type):
dtype = fe_type.set_type.dtype
members = [
# -1 = empty, -2 = deleted
('hash', types.intp),
('key', dtype),
]
super(SetEntryModel, self).__init__(dmm, fe_type, members)
@register_default(types.SetPayload)
class SetPayloadModel(StructModel):
def __init__(self, dmm, fe_type):
entry_type = types.SetEntry(fe_type.container)
members = [
# Number of active + deleted entries
('fill', types.intp),
# Number of active entries
('used', types.intp),
# Allocated size - 1 (size being a power of 2)
('mask', types.intp),
# Search finger
('finger', types.intp),
# This member is only used only for reflected sets
('dirty', types.boolean),
# Actually an inlined var-sized array
('entries', entry_type),
]
super(SetPayloadModel, self).__init__(dmm, fe_type, members)
@register_default(types.Set)
class SetModel(StructModel):
def __init__(self, dmm, fe_type):
payload_type = types.SetPayload(fe_type)
members = [
# The meminfo data points to a SetPayload
('meminfo', types.MemInfoPointer(payload_type)),
# This member is only used only for reflected sets
('parent', types.pyobject),
]
super(SetModel, self).__init__(dmm, fe_type, members)
@register_default(types.SetIter)
class SetIterModel(StructModel):
def __init__(self, dmm, fe_type):
payload_type = types.SetPayload(fe_type.container)
members = [
# The meminfo data points to a SetPayload (shared with the
# original set object)
('meminfo', types.MemInfoPointer(payload_type)),
# The index into the entries table
('index', types.EphemeralPointer(types.intp)),
]
super(SetIterModel, self).__init__(dmm, fe_type, members)
@register_default(types.Array)
@register_default(types.Buffer)
@register_default(types.ByteArray)
@register_default(types.Bytes)
@register_default(types.MemoryView)
@register_default(types.PyArray)
class ArrayModel(StructModel):
def __init__(self, dmm, fe_type):
ndim = fe_type.ndim
members = [
('meminfo', types.MemInfoPointer(fe_type.dtype)),
('parent', types.pyobject),
('nitems', types.intp),
('itemsize', types.intp),
('data', types.CPointer(fe_type.dtype)),
('shape', types.UniTuple(types.intp, ndim)),
('strides', types.UniTuple(types.intp, ndim)),
]
super(ArrayModel, self).__init__(dmm, fe_type, members)
@register_default(types.ArrayFlags)
class ArrayFlagsModel(StructModel):
def __init__(self, dmm, fe_type):
members = [
('parent', fe_type.array_type),
]
super(ArrayFlagsModel, self).__init__(dmm, fe_type, members)
@register_default(types.NestedArray)
class NestedArrayModel(ArrayModel):
def __init__(self, dmm, fe_type):
self._be_type = dmm.lookup(fe_type.dtype).get_data_type()
super(NestedArrayModel, self).__init__(dmm, fe_type)
@register_default(types.Optional)
class OptionalModel(StructModel):
def __init__(self, dmm, fe_type):
members = [
('data', fe_type.type),
('valid', types.boolean),
]
self._value_model = dmm.lookup(fe_type.type)
super(OptionalModel, self).__init__(dmm, fe_type, members)
def get_return_type(self):
return self._value_model.get_return_type()
def as_return(self, builder, value):
raise NotImplementedError
def from_return(self, builder, value):
return self._value_model.from_return(builder, value)
def traverse(self, builder):
def get_data(value):
valid = get_valid(value)
data = self.get(builder, value, "data")
return builder.select(valid, data, ir.Constant(data.type, None))
def get_valid(value):
return self.get(builder, value, "valid")
return [(self.get_type("data"), get_data),
(self.get_type("valid"), get_valid)]
@register_default(types.Record)
class RecordModel(CompositeModel):
def __init__(self, dmm, fe_type):
super(RecordModel, self).__init__(dmm, fe_type)
self._models = [self._dmm.lookup(t) for _, t in fe_type.members]
self._be_type = ir.ArrayType(ir.IntType(8), fe_type.size)
self._be_ptr_type = self._be_type.as_pointer()
def get_value_type(self):
"""Passed around as reference to underlying data
"""
return self._be_ptr_type
def get_argument_type(self):
return self._be_ptr_type
def get_return_type(self):
return self._be_ptr_type
def get_data_type(self):
return self._be_type
def as_data(self, builder, value):
return builder.load(value)
def from_data(self, builder, value):
raise NotImplementedError("use load_from_data_pointer() instead")
def as_argument(self, builder, value):
return value
def from_argument(self, builder, value):
return value
def as_return(self, builder, value):
return value
def from_return(self, builder, value):
return value
def load_from_data_pointer(self, builder, ptr, align=None):
return builder.bitcast(ptr, self.get_value_type())
@register_default(types.UnicodeCharSeq)
class UnicodeCharSeq(DataModel):
def __init__(self, dmm, fe_type):
super(UnicodeCharSeq, self).__init__(dmm, fe_type)
charty = ir.IntType(numpy_support.sizeof_unicode_char * 8)
self._be_type = ir.ArrayType(charty, fe_type.count)
def get_value_type(self):
return self._be_type
def get_data_type(self):
return self._be_type
def as_data(self, builder, value):
return value
def from_data(self, builder, value):
return value
def as_return(self, builder, value):
return value
def from_return(self, builder, value):
return value
def as_argument(self, builder, value):
return value
def from_argument(self, builder, value):
return value
@register_default(types.CharSeq)
class CharSeq(DataModel):
def __init__(self, dmm, fe_type):
super(CharSeq, self).__init__(dmm, fe_type)
charty = ir.IntType(8)
self._be_type = ir.ArrayType(charty, fe_type.count)
def get_value_type(self):
return self._be_type
def get_data_type(self):
return self._be_type
def as_data(self, builder, value):
return value
def from_data(self, builder, value):
return value
def as_return(self, builder, value):
return value
def from_return(self, builder, value):
return value
def as_argument(self, builder, value):
return value
def from_argument(self, builder, value):
return value
class CContiguousFlatIter(StructModel):
def __init__(self, dmm, fe_type, need_indices):
assert fe_type.array_type.layout == 'C'
array_type = fe_type.array_type
dtype = array_type.dtype
ndim = array_type.ndim
members = [('array', array_type),
('stride', types.intp),
('index', types.EphemeralPointer(types.intp)),
]
if need_indices:
# For ndenumerate()
members.append(('indices', types.EphemeralArray(types.intp, ndim)))
super(CContiguousFlatIter, self).__init__(dmm, fe_type, members)
class FlatIter(StructModel):
def __init__(self, dmm, fe_type):
array_type = fe_type.array_type
dtype = array_type.dtype
ndim = array_type.ndim
members = [('array', array_type),
('pointers', types.EphemeralArray(types.CPointer(dtype), ndim)),
('indices', types.EphemeralArray(types.intp, ndim)),
('exhausted', types.EphemeralPointer(types.boolean)),
]
super(FlatIter, self).__init__(dmm, fe_type, members)
@register_default(types.UniTupleIter)
class UniTupleIter(StructModel):
def __init__(self, dmm, fe_type):
members = [('index', types.EphemeralPointer(types.intp)),
('tuple', fe_type.container,)]
super(UniTupleIter, self).__init__(dmm, fe_type, members)
@register_default(types.misc.SliceLiteral)
@register_default(types.SliceType)
class SliceModel(StructModel):
def __init__(self, dmm, fe_type):
members = [('start', types.intp),
('stop', types.intp),
('step', types.intp),
]
super(SliceModel, self).__init__(dmm, fe_type, members)
@register_default(types.NPDatetime)
@register_default(types.NPTimedelta)
class NPDatetimeModel(PrimitiveModel):
def __init__(self, dmm, fe_type):
be_type = ir.IntType(64)
super(NPDatetimeModel, self).__init__(dmm, fe_type, be_type)
@register_default(types.ArrayIterator)
class ArrayIterator(StructModel):
def __init__(self, dmm, fe_type):
# We use an unsigned index to avoid the cost of negative index tests.
members = [('index', types.EphemeralPointer(types.uintp)),
('array', fe_type.array_type)]
super(ArrayIterator, self).__init__(dmm, fe_type, members)
@register_default(types.EnumerateType)
class EnumerateType(StructModel):
def __init__(self, dmm, fe_type):
members = [('count', types.EphemeralPointer(types.intp)),
('iter', fe_type.source_type)]
super(EnumerateType, self).__init__(dmm, fe_type, members)
@register_default(types.ZipType)
class ZipType(StructModel):
def __init__(self, dmm, fe_type):
members = [('iter%d' % i, source_type.iterator_type)
for i, source_type in enumerate(fe_type.source_types)]
super(ZipType, self).__init__(dmm, fe_type, members)
@register_default(types.RangeIteratorType)
class RangeIteratorType(StructModel):
def __init__(self, dmm, fe_type):
int_type = fe_type.yield_type
members = [('iter', types.EphemeralPointer(int_type)),
('stop', int_type),
('step', int_type),
('count', types.EphemeralPointer(int_type))]
super(RangeIteratorType, self).__init__(dmm, fe_type, members)
@register_default(types.Generator)
class GeneratorModel(CompositeModel):
def __init__(self, dmm, fe_type):
super(GeneratorModel, self).__init__(dmm, fe_type)
# XXX Fold this in DataPacker?
self._arg_models = [self._dmm.lookup(t) for t in fe_type.arg_types
if not isinstance(t, types.Omitted)]
self._state_models = [self._dmm.lookup(t) for t in fe_type.state_types]
self._args_be_type = ir.LiteralStructType(
[t.get_data_type() for t in self._arg_models])
self._state_be_type = ir.LiteralStructType(
[t.get_data_type() for t in self._state_models])
# The whole generator closure
self._be_type = ir.LiteralStructType(
[self._dmm.lookup(types.int32).get_value_type(),
self._args_be_type, self._state_be_type])
self._be_ptr_type = self._be_type.as_pointer()
def get_value_type(self):
"""
The generator closure is passed around as a reference.
"""
return self._be_ptr_type
def get_argument_type(self):
return self._be_ptr_type
def get_return_type(self):
return self._be_type
def get_data_type(self):
return self._be_type
def as_argument(self, builder, value):
return value
def from_argument(self, builder, value):
return value
def as_return(self, builder, value):
return self.as_data(builder, value)
def from_return(self, builder, value):
return self.from_data(builder, value)
def as_data(self, builder, value):
return builder.load(value)
def from_data(self, builder, value):
stack = cgutils.alloca_once(builder, value.type)
builder.store(value, stack)
return stack
@register_default(types.ArrayCTypes)
class ArrayCTypesModel(StructModel):
def __init__(self, dmm, fe_type):
# ndim = fe_type.ndim
members = [('data', types.CPointer(fe_type.dtype)),
('meminfo', types.MemInfoPointer(fe_type.dtype))]
super(ArrayCTypesModel, self).__init__(dmm, fe_type, members)
@register_default(types.RangeType)
class RangeModel(StructModel):
def __init__(self, dmm, fe_type):
int_type = fe_type.iterator_type.yield_type
members = [('start', int_type),
('stop', int_type),
('step', int_type)]
super(RangeModel, self).__init__(dmm, fe_type, members)
# =============================================================================
@register_default(types.NumpyNdIndexType)
class NdIndexModel(StructModel):
def __init__(self, dmm, fe_type):
ndim = fe_type.ndim
members = [('shape', types.UniTuple(types.intp, ndim)),
('indices', types.EphemeralArray(types.intp, ndim)),
('exhausted', types.EphemeralPointer(types.boolean)),
]
super(NdIndexModel, self).__init__(dmm, fe_type, members)
@register_default(types.NumpyFlatType)
def handle_numpy_flat_type(dmm, ty):
if ty.array_type.layout == 'C':
return CContiguousFlatIter(dmm, ty, need_indices=False)
else:
return FlatIter(dmm, ty)
@register_default(types.NumpyNdEnumerateType)
def handle_numpy_ndenumerate_type(dmm, ty):
if ty.array_type.layout == 'C':
return CContiguousFlatIter(dmm, ty, need_indices=True)
else:
return FlatIter(dmm, ty)
@register_default(types.BoundFunction)
def handle_bound_function(dmm, ty):
# The same as the underlying type
return dmm[ty.this]
@register_default(types.NumpyNdIterType)
class NdIter(StructModel):
def __init__(self, dmm, fe_type):
array_types = fe_type.arrays
ndim = fe_type.ndim
shape_len = ndim if fe_type.need_shaped_indexing else 1
members = [('exhausted', types.EphemeralPointer(types.boolean)),
('arrays', types.Tuple(array_types)),
# The iterator's main shape and indices
('shape', types.UniTuple(types.intp, shape_len)),
('indices', types.EphemeralArray(types.intp, shape_len)),
]
# Indexing state for the various sub-iterators
# XXX use a tuple instead?
for i, sub in enumerate(fe_type.indexers):
kind, start_dim, end_dim, _ = sub
member_name = 'index%d' % i
if kind == 'flat':
# A single index into the flattened array
members.append((member_name, types.EphemeralPointer(types.intp)))
elif kind in ('scalar', 'indexed', '0d'):
# Nothing required
pass
else:
assert 0
# Slots holding values of the scalar args
# XXX use a tuple instead?
for i, ty in enumerate(fe_type.arrays):
if not isinstance(ty, types.Array):
member_name = 'scalar%d' % i
members.append((member_name, types.EphemeralPointer(ty)))
super(NdIter, self).__init__(dmm, fe_type, members)
@register_default(types.DeferredType)
class DeferredStructModel(CompositeModel):
def __init__(self, dmm, fe_type):
super(DeferredStructModel, self).__init__(dmm, fe_type)
self.typename = "deferred.{0}".format(id(fe_type))
self.actual_fe_type = fe_type.get()
def get_value_type(self):
return ir.global_context.get_identified_type(self.typename + '.value')
def get_data_type(self):
return ir.global_context.get_identified_type(self.typename + '.data')
def get_argument_type(self):
return self._actual_model.get_argument_type()
def as_argument(self, builder, value):
inner = self.get(builder, value)
return self._actual_model.as_argument(builder, inner)
def from_argument(self, builder, value):
res = self._actual_model.from_argument(builder, value)
return self.set(builder, self.make_uninitialized(), res)
def from_data(self, builder, value):
self._define()
elem = self.get(builder, value)
value = self._actual_model.from_data(builder, elem)
out = self.make_uninitialized()
return self.set(builder, out, value)
def as_data(self, builder, value):
self._define()
elem = self.get(builder, value)
value = self._actual_model.as_data(builder, elem)
out = self.make_uninitialized(kind='data')
return self.set(builder, out, value)
def from_return(self, builder, value):
return value
def as_return(self, builder, value):
return value
def get(self, builder, value):
return builder.extract_value(value, [0])
def set(self, builder, value, content):
return builder.insert_value(value, content, [0])
def make_uninitialized(self, kind='value'):
self._define()
if kind == 'value':
ty = self.get_value_type()
else:
ty = self.get_data_type()
return ir.Constant(ty, ir.Undefined)
def _define(self):
valty = self.get_value_type()
self._define_value_type(valty)
datty = self.get_data_type()
self._define_data_type(datty)
def _define_value_type(self, value_type):
if value_type.is_opaque:
value_type.set_body(self._actual_model.get_value_type())
def _define_data_type(self, data_type):
if data_type.is_opaque:
data_type.set_body(self._actual_model.get_data_type())
@property
def _actual_model(self):
return self._dmm.lookup(self.actual_fe_type)
def traverse(self, builder):
return [(self.actual_fe_type,
lambda value: builder.extract_value(value, [0]))]
@register_default(types.StructRefPayload)
class StructPayloadModel(StructModel):
"""Model for the payload of a mutable struct
"""
def __init__(self, dmm, fe_typ):
members = tuple(fe_typ.field_dict.items())
super().__init__(dmm, fe_typ, members)
class StructRefModel(StructModel):
"""Model for a mutable struct.
A reference to the payload
"""
def __init__(self, dmm, fe_typ):
dtype = fe_typ.get_data_type()
members = [
("meminfo", types.MemInfoPointer(dtype)),
]
super().__init__(dmm, fe_typ, members)