def make_vector_type(name: str, base_type: types.Type, attr_names: Tuple[str, ...], user_facing_object) -> types.Type: """Create a vector type. Parameters ---------- name: str The name of the type. base_type: numba.types.Type The primitive type for each element in the vector. attr_names: tuple of str Name for each attribute. user_facing_object: object The handle to be used in cuda kernel. """ class _VectorType(VectorType): """Internal instantiation of VectorType.""" pass class VectorTypeModel(models.StructModel): def __init__(self, dmm, fe_type): members = [(attr_name, base_type) for attr_name in attr_names] super().__init__(dmm, fe_type, members) vector_type = _VectorType(name, base_type, attr_names, user_facing_object) register_model(_VectorType)(VectorTypeModel) for attr_name in attr_names: make_attribute_wrapper(_VectorType, attr_name, attr_name) return vector_type
def setUp(self): class Dummy(object): def __init__(self, value): self.value = value class DummyType(types.Type): def __init__(self): super(DummyType, self).__init__(name="Dummy") dummy_type = DummyType() @register_model(DummyType) class DummyModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ("value", types.intp), ] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(DummyType, "value", "value") @type_callable(Dummy) def type_dummy(context): def typer(value): return dummy_type return typer @lower_builtin(Dummy, types.intp) def impl_dummy(context, builder, sig, args): typ = sig.return_type [value] = args dummy = cgutils.create_struct_proxy(typ)(context, builder) dummy.value = value return dummy._getvalue() @typeof_impl.register(Dummy) def typeof_dummy(val, c): return DummyType() # Store attributes self.Dummy = Dummy self.DummyType = DummyType
def test_issue5792(self): # Issue is that overloads cache their IR and closure inliner was # manipulating the cached IR in a way that broke repeated inlines. class Dummy: def __init__(self, data): self.data = data def div(self, other): return data / other.data class DummyType(types.Type): def __init__(self, data): self.data = data super().__init__(name=f'Dummy({self.data})') @register_model(DummyType) class DummyTypeModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ('data', fe_type.data), ] super().__init__(dmm, fe_type, members) make_attribute_wrapper(DummyType, 'data', '_data') @intrinsic def init_dummy(typingctx, data): def codegen(context, builder, sig, args): typ = sig.return_type data, = args dummy = cgutils.create_struct_proxy(typ)(context, builder) dummy.data = data if context.enable_nrt: context.nrt.incref(builder, sig.args[0], data) return dummy._getvalue() ret_typ = DummyType(data) sig = signature(ret_typ, data) return sig, codegen @overload(Dummy, inline='always') def dummy_overload(data): def ctor(data): return init_dummy(data) return ctor @overload_method(DummyType, 'div', inline='always') def div_overload(self, other): def impl(self, other): return self._data / other._data return impl @njit def test_impl(data, other_data): dummy = Dummy(data) # ctor inlined once other = Dummy(other_data) # ctor inlined again return dummy.div(other) data = 1. other_data = 2. res = test_impl(data, other_data) self.assertEqual(res, data / other_data)
self).__init__(name='IndexValueType({})'.format(val_typ)) @typeof_impl.register(IndexValue) def typeof_index(val, c): val_typ = typeof_impl(val.value, c) return IndexValueType(val_typ) @type_callable(IndexValue) def type_index_value(context): def typer(ind, mval): if ind == types.intp or ind == types.uintp: return IndexValueType(mval) return typer @register_model(IndexValueType) class IndexValueModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ('index', types.intp), ('value', fe_type.val_typ), ] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(IndexValueType, 'index', 'index') make_attribute_wrapper(IndexValueType, 'value', 'value')
@register_model(DictItemsIterableType) @register_model(DictKeysIterableType) @register_model(DictValuesIterableType) @register_model(DictIteratorType) class DictIterModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ('parent', fe_type.parent), # reference to the dict ('state', types.voidptr), # iterator state in C code ] super(DictIterModel, self).__init__(dmm, fe_type, members) # Make _parent available to make len simple make_attribute_wrapper(DictItemsIterableType, "parent", "_parent") make_attribute_wrapper(DictKeysIterableType, "parent", "_parent") make_attribute_wrapper(DictValuesIterableType, "parent", "_parent") def _raise_if_error(context, builder, status, msg): """Raise an internal error depending on the value of *status* """ ok_status = status.type(int(Status.OK)) with builder.if_then(builder.icmp_signed('!=', status, ok_status)): context.call_conv.return_user_exc(builder, RuntimeError, (msg, )) @intrinsic def _as_meminfo(typingctx, dctobj): """Returns the MemInfoPointer of a dictionary.
('hi', types.float64), ] models.StructModel.__init__(self, dmm, fe_type, members) @lower_builtin(Interval, types.Float, types.Float) def impl_interval(context, builder, sig, args): typ = sig.return_type lo, hi = args interval = cgutils.create_struct_proxy(typ)(context, builder) interval.lo = lo interval.hi = hi return interval._getvalue() make_attribute_wrapper(IntervalType, 'lo', 'lo') make_attribute_wrapper(IntervalType, 'hi', 'hi') # From the tutorial - doesn't work due to: # # No definition for lowering <built-in method getter of _dynfunc._Closure # object at 0x7ff6d564c4c0>(Interval,) -> float64 # @overload_attribute(IntervalType, "width") # def get_width(interval): # def getter(interval): # return interval.hi - interval.lo # return getter # Alternative:
datetime_cases = {types.NPDatetime(u) for u in units} timedelta_cases = {types.NPTimedelta(u) for u in units} cases = [ nb_signature(MaskedType(t), t, types.boolean) for t in ( types.integer_domain | types.real_domain | datetime_cases | timedelta_cases | {types.boolean} ) ] # Provide access to `m.value` and `m.valid` in a kernel for a Masked `m`. make_attribute_wrapper(MaskedType, "value", "value") make_attribute_wrapper(MaskedType, "valid", "valid") # Typing for `api.Masked` @cuda_decl_registry.register_attr class ClassesTemplate(AttributeTemplate): key = types.Module(api) def resolve_Masked(self, mod): return types.Function(MaskedConstructor) # Registration of the global is also needed for Numba to type api.Masked cuda_decl_registry.register_global(api, types.Module(api)) # For typing bare Masked (as in `from .api import Masked`
def __init__(self, dmm, fe_type): members = [('data', fe_type.as_array)] models.StructModel.__init__(self, dmm, fe_type, members) @register_model(SeriesType) class SeriesModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ('index', fe_type.index), ('values', fe_type.as_array), ] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(IndexType, 'data', '_data') make_attribute_wrapper(SeriesType, 'index', '_index') make_attribute_wrapper(SeriesType, 'values', '_values') def make_index(context, builder, typ, **kwargs): return cgutils.create_struct_proxy(typ)(context, builder, **kwargs) def make_series(context, builder, typ, **kwargs): return cgutils.create_struct_proxy(typ)(context, builder, **kwargs) @lower_builtin('__array__', IndexType) def index_as_array(context, builder, sig, args): val = make_index(context, builder, sig.args[0], ref=args[0])
self).__init__(name="IndexValueType({})".format(val_typ)) @typeof_impl.register(IndexValue) def typeof_index(val, c): val_typ = typeof_impl(val.value, c) return IndexValueType(val_typ) @type_callable(IndexValue) def type_index_value(context): def typer(ind, mval): if ind == types.intp or ind == types.uintp: return IndexValueType(mval) return typer @register_model(IndexValueType) class IndexValueModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ("index", types.intp), ("value", fe_type.val_typ), ] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(IndexValueType, "index", "index") make_attribute_wrapper(IndexValueType, "value", "value")
def __init__(self, dmm, fe_type): members = [("data", fe_type.as_array)] models.StructModel.__init__(self, dmm, fe_type, members) @register_model(SeriesType) class SeriesModel(models.StructModel): def __init__(self, dmm, fe_type): members = [ ("index", fe_type.index), ("values", fe_type.as_array), ] models.StructModel.__init__(self, dmm, fe_type, members) make_attribute_wrapper(IndexType, "data", "_data") make_attribute_wrapper(SeriesType, "index", "_index") make_attribute_wrapper(SeriesType, "values", "_values") def make_index(context, builder, typ, **kwargs): return cgutils.create_struct_proxy(typ)(context, builder, **kwargs) def make_series(context, builder, typ, **kwargs): return cgutils.create_struct_proxy(typ)(context, builder, **kwargs) @lower_builtin("__array__", IndexType) def index_as_array(context, builder, sig, args): val = make_index(context, builder, sig.args[0], ref=args[0])
models.StructModel.__init__(self, dmm, fe_type, members) @lower_builtin(Quaternion, types.Float, types.Float, types.Float, types.Float) def impl_quaternion(context, builder, sig, args): typ = sig.return_type a, b, c, d = args quaternion = cgutils.create_struct_proxy(typ)(context, builder) quaternion.a = a quaternion.b = b quaternion.c = c quaternion.d = d return quaternion._getvalue() make_attribute_wrapper(QuaternionType, 'a', 'a') make_attribute_wrapper(QuaternionType, 'b', 'b') make_attribute_wrapper(QuaternionType, 'c', 'c') make_attribute_wrapper(QuaternionType, 'd', 'd') @cuda_registry.register_attr class Quaternion_attrs(AttributeTemplate): key = QuaternionType def resolve_phi(self, mod): return types.float64 def resolve_theta(self, mod): return types.float64
make_attribute_wrapper, typeof_impl, type_callable) from numba.cuda.cudaimpl import lower from numba.core import cgutils @typeof_impl.register(TestStruct) def typeof_teststruct(val, c): return test_struct_model_type @register_model(TestStructModelType) class TestStructModel(models.StructModel): def __init__(self, dmm, fe_type): members = [("x", int32), ("y", int32)] super().__init__(dmm, fe_type, members) make_attribute_wrapper(TestStructModelType, 'x', 'x') make_attribute_wrapper(TestStructModelType, 'y', 'y') @type_callable(TestStruct) def type_test_struct(context): def typer(x, y): if isinstance(x, types.Integer) and isinstance(y, types.Integer): return test_struct_model_type return typer @lower(TestStruct, types.Integer, types.Integer) def lower_test_type_ctor(context, builder, sig, args): obj = cgutils.create_struct_proxy(test_struct_model_type)(context, builder) obj.x = args[0]
_bit_gen_type = types.NumPyRandomBitGeneratorType('bit_generator') @register_model(types.NumPyRandomGeneratorType) class NumPyRandomGeneratorTypeModel(models.StructModel): def __init__(self, dmm, fe_type): members = [('bit_generator', _bit_gen_type), ('meminfo', types.MemInfoPointer(types.voidptr)), ('parent', types.pyobject)] super(NumPyRandomGeneratorTypeModel, self).__init__(dmm, fe_type, members) # The Generator instances have a bit_generator attr make_attribute_wrapper(types.NumPyRandomGeneratorType, 'bit_generator', 'bit_generator') def _generate_next_binding(overloadable_function, return_type): """ Generate the overloads for "next_(some type)" functions. """ @intrinsic def intrin_NumPyRandomBitGeneratorType_next_ty(tyctx, inst): sig = return_type(inst) def codegen(cgctx, builder, sig, llargs): name = overloadable_function.__name__ struct_ptr = cgutils.create_struct_proxy(inst)(cgctx, builder, value=llargs[0])