def _getsetitem_gen(getset):
_dunder_meth = "__%s__" % getset
op = getattr(operator, getset)
@templates.infer_global(op)
class GetSetItem(templates.AbstractTemplate):
def generic(self, args, kws):
instance = args[0]
if (
isinstance(instance, types.ClassInstanceType)
and _dunder_meth in instance.jitmethods
):
meth = instance.jitmethods[_dunder_meth]
disp_type = types.Dispatcher(meth)
sig = disp_type.get_call_type(self.context, args, kws)
return sig
# lower both {g,s}etitem and __{g,s}etitem__ to catch the calls
# from python and numba
imputils.lower_builtin(
(types.ClassInstanceType, _dunder_meth),
types.ClassInstanceType,
types.VarArg(types.Any),
)(get_imp())
imputils.lower_builtin(
op, types.ClassInstanceType, types.VarArg(types.Any)
)(get_imp())
def generic(self, args, kws):
"""
Type the intrinsic by the arguments.
"""
from numba.core.extending_hardware import resolve_dispatcher_from_str
from numba.core.imputils import builtin_registry
cache_key = self.context, args, tuple(kws.items())
hwstr = self.metadata.get('hardware', 'cpu')
disp = resolve_dispatcher_from_str(hwstr)
tgtctx = disp.targetdescr.target_context
# This is all workarounds...
# The issue is that whilst targets shouldn't care about which registry
# in which to register lowering implementations, the CUDA target
# "borrows" implementations from the CPU from specific registries. This
# means that if some impl is defined via @intrinsic, e.g. numba.*unsafe
# modules, _AND_ CUDA also makes use of the same impl, then it's
# required that the registry in use is one that CUDA borrows from. This
# leads to the following expression where by the CPU builtin_registry is
# used if it is in the target context as a known registry (i.e. the
# target installed it) and if it is not then it is assumed that the
# registries for the target are unbound to any other target and so it's
# fine to use any of them as a place to put lowering impls.
#
# NOTE: This will need subsequently fixing again when targets use solely
# the extension APIs to describe their implementation. The issue will be
# that the builtin_registry should contain _just_ the stack allocated
# implementations and low level target invariant things and should not
# be modified further. It should be acceptable to remove the `then`
# branch and just keep the `else`.
# In case the target has swapped, e.g. cuda borrowing cpu, refresh to
# populate.
tgtctx.refresh()
if builtin_registry in tgtctx._registries:
reg = builtin_registry
else:
# Pick a registry in which to install intrinsics
registries = iter(tgtctx._registries)
reg = next(registries)
lower_builtin = reg.lower
try:
return self._impl_cache[cache_key]
except KeyError:
pass
result = self._definition_func(self.context, *args, **kws)
if result is None:
return
[sig, imp] = result
pysig = utils.pysignature(self._definition_func)
# omit context argument from user function
parameters = list(pysig.parameters.values())[1:]
sig = sig.replace(pysig=pysig.replace(parameters=parameters))
self._impl_cache[cache_key] = sig
self._overload_cache[sig.args] = imp
# register the lowering
lower_builtin(imp, *sig.args)(imp)
return sig
def impl_ldexp():
ldexp_cpu = gen_codegen('ldexp')
ldexp_gpu = gen_codegen('__nv_ldexp')
ldexpf_cpu = gen_codegen('ldexpf')
ldexpf_gpu = gen_codegen('__nv_ldexpf')
lower_builtin(math.ldexp, float64,
int32)(dispatch_codegen(ldexp_cpu, ldexp_gpu))
lower_builtin(math.ldexp, float32,
int32)(dispatch_codegen(ldexpf_cpu, ldexpf_gpu))
def generic(self, args, kws):
"""
Type the intrinsic by the arguments.
"""
from numba.core.imputils import lower_builtin
cache_key = self.context, args, tuple(kws.items())
try:
return self._impl_cache[cache_key]
except KeyError:
result = self._definition_func(self.context, *args, **kws)
if result is None:
return
[sig, imp] = result
pysig = utils.pysignature(self._definition_func)
# omit context argument from user function
parameters = list(pysig.parameters.values())[1:]
sig = sig.replace(pysig=pysig.replace(parameters=parameters))
self._impl_cache[cache_key] = sig
self._overload_cache[sig.args] = imp
# register the lowering
lower_builtin(imp, *sig.args)(imp)
return sig
def _implement_bitwise_operators():
for ty in (types.Boolean, types.Integer):
lower_builtin(operator.and_, ty, ty)(int_and_impl)
lower_builtin(operator.iand, ty, ty)(int_and_impl)
lower_builtin(operator.or_, ty, ty)(int_or_impl)
lower_builtin(operator.ior, ty, ty)(int_or_impl)
lower_builtin(operator.xor, ty, ty)(int_xor_impl)
lower_builtin(operator.ixor, ty, ty)(int_xor_impl)
lower_builtin(operator.invert, ty)(int_invert_impl)
def _implement_integer_operators():
ty = types.Integer
lower_builtin(operator.add, ty, ty)(int_add_impl)
lower_builtin(operator.iadd, ty, ty)(int_add_impl)
lower_builtin(operator.sub, ty, ty)(int_sub_impl)
lower_builtin(operator.isub, ty, ty)(int_sub_impl)
lower_builtin(operator.mul, ty, ty)(int_mul_impl)
lower_builtin(operator.imul, ty, ty)(int_mul_impl)
lower_builtin(operator.eq, ty, ty)(int_eq_impl)
lower_builtin(operator.ne, ty, ty)(int_ne_impl)
lower_builtin(operator.lshift, ty, ty)(int_shl_impl)
lower_builtin(operator.ilshift, ty, ty)(int_shl_impl)
lower_builtin(operator.rshift, ty, ty)(int_shr_impl)
lower_builtin(operator.irshift, ty, ty)(int_shr_impl)
lower_builtin(operator.neg, ty)(int_negate_impl)
lower_builtin(operator.pos, ty)(int_positive_impl)
lower_builtin(operator.pow, ty, ty)(int_power_impl)
lower_builtin(operator.ipow, ty, ty)(int_power_impl)
lower_builtin(pow, ty, ty)(int_power_impl)
for ty in types.unsigned_domain:
lower_builtin(operator.lt, ty, ty)(int_ult_impl)
lower_builtin(operator.le, ty, ty)(int_ule_impl)
lower_builtin(operator.gt, ty, ty)(int_ugt_impl)
lower_builtin(operator.ge, ty, ty)(int_uge_impl)
lower_builtin(operator.pow, types.Float, ty)(int_power_impl)
lower_builtin(operator.ipow, types.Float, ty)(int_power_impl)
lower_builtin(pow, types.Float, ty)(int_power_impl)
lower_builtin(abs, ty)(uint_abs_impl)
lower_builtin(operator.lt, types.IntegerLiteral,
types.IntegerLiteral)(int_slt_impl)
lower_builtin(operator.gt, types.IntegerLiteral,
types.IntegerLiteral)(int_slt_impl)
lower_builtin(operator.le, types.IntegerLiteral,
types.IntegerLiteral)(int_slt_impl)
lower_builtin(operator.ge, types.IntegerLiteral,
types.IntegerLiteral)(int_slt_impl)
for ty in types.signed_domain:
lower_builtin(operator.lt, ty, ty)(int_slt_impl)
lower_builtin(operator.le, ty, ty)(int_sle_impl)
lower_builtin(operator.gt, ty, ty)(int_sgt_impl)
lower_builtin(operator.ge, ty, ty)(int_sge_impl)
lower_builtin(operator.pow, types.Float, ty)(int_power_impl)
lower_builtin(operator.ipow, types.Float, ty)(int_power_impl)
lower_builtin(pow, types.Float, ty)(int_power_impl)
lower_builtin(abs, ty)(int_abs_impl)
def bool_negate_impl(context, builder, sig, args):
[typ] = sig.args
[val] = args
res = context.cast(builder, val, typ, sig.return_type)
res = builder.neg(res)
return impl_ret_untracked(context, builder, sig.return_type, res)
def bool_unary_positive_impl(context, builder, sig, args):
[typ] = sig.args
[val] = args
res = context.cast(builder, val, typ, sig.return_type)
return impl_ret_untracked(context, builder, sig.return_type, res)
lower_builtin(operator.eq, types.boolean, types.boolean)(int_eq_impl)
lower_builtin(operator.ne, types.boolean, types.boolean)(int_ne_impl)
lower_builtin(operator.lt, types.boolean, types.boolean)(int_ult_impl)
lower_builtin(operator.le, types.boolean, types.boolean)(int_ule_impl)
lower_builtin(operator.gt, types.boolean, types.boolean)(int_ugt_impl)
lower_builtin(operator.ge, types.boolean, types.boolean)(int_uge_impl)
lower_builtin(operator.neg, types.boolean)(bool_negate_impl)
lower_builtin(operator.pos, types.boolean)(bool_unary_positive_impl)
def _implement_integer_operators():
ty = types.Integer
lower_builtin(operator.add, ty, ty)(int_add_impl)
lower_builtin(operator.iadd, ty, ty)(int_add_impl)
lower_builtin(operator.sub, ty, ty)(int_sub_impl)
types.Literal.ctor_map[bool] = BooleanLiteral
register_default(BooleanLiteral)(numba.core.datamodel.models.BooleanModel)
@lower_cast(BooleanLiteral, types.boolean)
def literal_bool_cast(context, builder, fromty, toty, val):
lit = context.get_constant_generic(
builder,
fromty.literal_type,
fromty.literal_value,
)
return context.cast(builder, lit, fromty.literal_type, toty)
lower_builtin(operator.eq, BooleanLiteral,
BooleanLiteral)(numba.cpython.builtins.const_eq_impl)
lower_builtin(operator.ne, BooleanLiteral,
BooleanLiteral)(numba.cpython.builtins.const_ne_impl)
@lower_builtin(bool, BooleanLiteral)
def bool_as_bool(context, builder, sig, args):
[val] = args
return val
def get_constant(func_ir, var, default=NOT_CONSTANT):
def_node = guard(get_definition, func_ir, var)
if def_node is None:
return default
if isinstance(def_node, ir.Const):
timedelta_eq_timedelta_impl = _create_timedelta_comparison_impl(
'==', cgutils.false_bit)
timedelta_ne_timedelta_impl = _create_timedelta_comparison_impl(
'!=', cgutils.true_bit)
timedelta_lt_timedelta_impl = _create_timedelta_ordering_impl('<')
timedelta_le_timedelta_impl = _create_timedelta_ordering_impl('<=')
timedelta_gt_timedelta_impl = _create_timedelta_ordering_impl('>')
timedelta_ge_timedelta_impl = _create_timedelta_ordering_impl('>=')
for op_, func in [(operator.eq, timedelta_eq_timedelta_impl),
(operator.ne, timedelta_ne_timedelta_impl),
(operator.lt, timedelta_lt_timedelta_impl),
(operator.le, timedelta_le_timedelta_impl),
(operator.gt, timedelta_gt_timedelta_impl),
(operator.ge, timedelta_ge_timedelta_impl)]:
lower_builtin(op_, *TIMEDELTA_BINOP_SIG)(func)
# Arithmetic on datetime64
def is_leap_year(builder, year_val):
"""
Return a predicate indicating whether *year_val* (offset by 1970) is a
leap year.
"""
actual_year = builder.add(year_val, Constant(DATETIME64, 1970))
multiple_of_4 = cgutils.is_null(
builder, builder.and_(actual_year, Constant(DATETIME64, 3)))
not_multiple_of_100 = cgutils.is_not_null(
builder, builder.srem(actual_year, Constant(DATETIME64, 100)))
multiple_of_400 = cgutils.is_null(
# Make sure None is on the right
if lty == types.none:
lty, rty = rty, lty
lval, rval = rval, lval
opt_type = lty
opt_val = lval
opt = context.make_helper(builder, opt_type, opt_val)
res = builder.not_(cgutils.as_bool_bit(builder, opt.valid))
return impl_ret_untracked(context, builder, sig.return_type, res)
# None is/not None
lower_builtin(operator.is_, types.none, types.none)(always_return_true_impl)
# Optional is None
lower_builtin(operator.is_, types.Optional, types.none)(optional_is_none)
lower_builtin(operator.is_, types.none, types.Optional)(optional_is_none)
@lower_getattr_generic(types.Optional)
def optional_getattr(context, builder, typ, value, attr):
"""
Optional.__getattr__ => redirect to the wrapped type.
"""
inner_type = typ.type
val = context.cast(builder, value, typ, inner_type)
imp = context.get_getattr(inner_type, attr)
return imp(context, builder, inner_type, val, attr)
def impl_binary(fname, key, typ):
cpu = gen_codegen(fname)
gpu = gen_codegen(f"__nv_{fname}")
lower_builtin(key, typ, typ)(dispatch_codegen(cpu, gpu))
