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)