def __init__(self, context, builder, outer_sig):
super(_KernelImpl, self).__init__(context, builder, outer_sig)
loop = ufunc_find_matching_loop(
ufunc, outer_sig.args + tuple(_unpack_output_types(ufunc, outer_sig)))
self.fn = context.get_ufunc_info(ufunc).get(loop.ufunc_sig)
self.inner_sig = _ufunc_loop_sig(loop.outputs, loop.inputs)
if self.fn is None:
msg = "Don't know how to lower ufunc '{0}' for loop '{1}'"
raise NotImplementedError(msg.format(ufunc.__name__, loop))
def generic(self, args, kws):
# First, strip optional types, ufunc loops are typed on concrete types
args = [x.type if isinstance(x, types.Optional) else x for x in args]
ufunc = self.ufunc
base_types, explicit_outputs, ndims, layout = self._handle_inputs(
ufunc, args, kws)
ufunc_loop = ufunc_find_matching_loop(ufunc, base_types)
if ufunc_loop is None:
raise TypingError("can't resolve ufunc {0} for types {1}".format(
ufunc.__name__, args))
# check if all the types involved in the ufunc loop are supported in this mode
if not supported_ufunc_loop(ufunc, ufunc_loop):
msg = "ufunc '{0}' using the loop '{1}' not supported in this mode"
raise TypingError(
msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig))
# if there is any explicit output type, check that it is valid
explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs]
# Numpy will happily use unsafe conversions (although it will actually warn)
if not all(
np.can_cast(fromty, toty, 'unsafe')
for (fromty, toty
) in zip(ufunc_loop.numpy_outputs, explicit_outputs_np)):
msg = "ufunc '{0}' can't cast result to explicit result type"
raise TypingError(msg=msg.format(ufunc.__name__))
# A valid loop was found that is compatible. The result of type inference should
# be based on the explicit output types, and when not available with the type given
# by the selected NumPy loop
out = list(explicit_outputs)
implicit_output_count = ufunc.nout - len(explicit_outputs)
if implicit_output_count > 0:
# XXX this is sometimes wrong for datetime64 and timedelta64,
# as ufunc_find_matching_loop() doesn't do any type inference
ret_tys = ufunc_loop.outputs[-implicit_output_count:]
if ndims > 0:
assert layout is not None
ret_tys = [
types.Array(dtype=ret_ty, ndim=ndims, layout=layout)
for ret_ty in ret_tys
]
ret_tys = [
resolve_output_type(self.context, args, ret_ty)
for ret_ty in ret_tys
]
out.extend(ret_tys)
return _ufunc_loop_sig(out, args)
def numpy_ufunc_kernel(context, builder, sig, args, ufunc, kernel_class):
# This is the code generator that builds all the looping needed
# to execute a numpy functions over several dimensions (including
# scalar cases).
#
# context - the code generation context
# builder - the code emitter
# sig - signature of the ufunc
# args - the args to the ufunc
# ufunc - the ufunc itself
# kernel_class - a code generating subclass of _Kernel that provides
arguments = [
_prepare_argument(context, builder, arg, tyarg)
for arg, tyarg in zip(args, sig.args)
]
if len(arguments) < ufunc.nin:
raise RuntimeError(
"Not enough inputs to {}, expected {} got {}".format(
ufunc.__name__, ufunc.nin, len(arguments)))
for out_i, ret_ty in enumerate(_unpack_output_types(ufunc, sig)):
if ufunc.nin + out_i >= len(arguments):
# this out argument is not provided
if isinstance(ret_ty, types.ArrayCompatible):
output = _build_array(context, builder, ret_ty, sig.args,
arguments)
else:
output = _prepare_argument(
context, builder,
lc.Constant.null(context.get_value_type(ret_ty)), ret_ty)
arguments.append(output)
elif context.enable_nrt:
# Incref the output
context.nrt.incref(builder, ret_ty, args[ufunc.nin + out_i])
inputs = arguments[:ufunc.nin]
outputs = arguments[ufunc.nin:]
assert len(outputs) == ufunc.nout
outer_sig = _ufunc_loop_sig([a.base_type for a in outputs],
[a.base_type for a in inputs])
kernel = kernel_class(context, builder, outer_sig)
intpty = context.get_value_type(types.intp)
indices = [inp.create_iter_indices() for inp in inputs]
# assume outputs are all the same size, which numpy requires
loopshape = outputs[0].shape
with cgutils.loop_nest(builder, loopshape, intp=intpty) as loop_indices:
vals_in = []
for i, (index, arg) in enumerate(zip(indices, inputs)):
index.update_indices(loop_indices, i)
vals_in.append(arg.load_data(index.as_values()))
vals_out = _unpack_output_values(ufunc, builder,
kernel.generate(*vals_in))
for val_out, output in zip(vals_out, outputs):
output.store_data(loop_indices, val_out)
out = _pack_output_values(ufunc, context, builder, sig.return_type,
[o.return_val for o in outputs])
return impl_ret_new_ref(context, builder, sig.return_type, out)
def generic(self, args, kws):
# First, strip optional types, ufunc loops are typed on concrete types
args = [x.type if isinstance(x, types.Optional) else x for x in args]
ufunc = self.ufunc
base_types, explicit_outputs, ndims, layout = self._handle_inputs(
ufunc, args, kws)
ufunc_loop = ufunc_find_matching_loop(ufunc, base_types)
if ufunc_loop is None:
raise TypingError("can't resolve ufunc {0} for types {1}".format(
ufunc.__name__, args))
# check if all the types involved in the ufunc loop are supported in this mode
if not supported_ufunc_loop(ufunc, ufunc_loop):
msg = "ufunc '{0}' using the loop '{1}' not supported in this mode"
raise TypingError(
msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig))
# if there is any explicit output type, check that it is valid
explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs]
# Numpy will happily use unsafe conversions (although it will actually warn)
if not all(
np.can_cast(fromty, toty, 'unsafe')
for (fromty, toty
) in zip(ufunc_loop.numpy_outputs, explicit_outputs_np)):
msg = "ufunc '{0}' can't cast result to explicit result type"
raise TypingError(msg=msg.format(ufunc.__name__))
# A valid loop was found that is compatible. The result of type inference should
# be based on the explicit output types, and when not available with the type given
# by the selected NumPy loop
out = list(explicit_outputs)
implicit_output_count = ufunc.nout - len(explicit_outputs)
if implicit_output_count > 0:
# XXX this is sometimes wrong for datetime64 and timedelta64,
# as ufunc_find_matching_loop() doesn't do any type inference
ret_tys = ufunc_loop.outputs[-implicit_output_count:]
if ndims > 0:
assert layout is not None
# If either of the types involved in the ufunc operation have a
# __array_ufunc__ method then invoke the first such one to
# determine the output type of the ufunc.
array_ufunc_type = None
for a in args:
if hasattr(a, "__array_ufunc__"):
array_ufunc_type = a
break
output_type = types.Array
if array_ufunc_type is not None:
output_type = array_ufunc_type.__array_ufunc__(
ufunc, "__call__", *args, **kws)
if output_type is NotImplemented:
msg = (f"unsupported use of ufunc {ufunc} on "
f"{array_ufunc_type}")
# raise TypeError here because
# NumpyRulesArrayOperator.generic is capturing
# TypingError
raise NumbaTypeError(msg)
elif not issubclass(output_type, types.Array):
msg = (f"ufunc {ufunc} on {array_ufunc_type}"
f"cannot return non-array {output_type}")
# raise TypeError here because
# NumpyRulesArrayOperator.generic is capturing
# TypingError
raise TypeError(msg)
ret_tys = [
output_type(dtype=ret_ty, ndim=ndims, layout=layout)
for ret_ty in ret_tys
]
ret_tys = [
resolve_output_type(self.context, args, ret_ty)
for ret_ty in ret_tys
]
out.extend(ret_tys)
return _ufunc_loop_sig(out, args)
