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))