""" Implements custom ufunc dispatch mechanism for non-CPU devices. """ from abc import ABCMeta, abstractmethod from collections import OrderedDict import operator import warnings from functools import reduce import numpy as np from numba.np.ufunc.ufuncbuilder import _BaseUFuncBuilder, parse_identity from numba.core import types, sigutils from numba.core.typing import signature from numba.np.ufunc.sigparse import parse_signature def _broadcast_axis(a, b): """ Raises ------ ValueError if broadcast fails """ if a == b: return a elif a == 1: return b elif b == 1: return a else: raise ValueError("failed to broadcast {0} and {1}".format(a, b)) def _pairwise_broadcast(shape1, shape2): """ Raises ------ ValueError if broadcast fails """ shape1, shape2 = map(tuple, [shape1, shape2]) while len(shape1) < len(shape2): shape1 = (1,) + shape1 while len(shape1) > len(shape2): shape2 = (1,) + shape2 return tuple(_broadcast_axis(a, b) for a, b in zip(shape1, shape2)) def _multi_broadcast(*shapelist): """ Raises ------ ValueError if broadcast fails """ assert shapelist result = shapelist[0] others = shapelist[1:] try: for i, each in enumerate(others, start=1): result = _pairwise_broadcast(result, each) except ValueError: raise ValueError("failed to broadcast argument #{0}".format(i)) else: return result class UFuncMechanism(object): """ Prepare ufunc arguments for vectorize. """ DEFAULT_STREAM = None SUPPORT_DEVICE_SLICING = False def __init__(self, typemap, args): """Never used directly by user. Invoke by UFuncMechanism.call(). """ self.typemap = typemap self.args = args nargs = len(self.args) self.argtypes = [None] * nargs self.scalarpos = [] self.signature = None self.arrays = [None] * nargs def _fill_arrays(self): """ Get all arguments in array form """ for i, arg in enumerate(self.args): if self.is_device_array(arg): self.arrays[i] = self.as_device_array(arg) elif isinstance(arg, (int, float, complex, np.number)): # Is scalar self.scalarpos.append(i) else: self.arrays[i] = np.asarray(arg) def _fill_argtypes(self): """ Get dtypes """ for i, ary in enumerate(self.arrays): if ary is not None: dtype = getattr(ary, 'dtype') if dtype is None: dtype = np.asarray(ary).dtype self.argtypes[i] = dtype def _resolve_signature(self): """Resolve signature. May have ambiguous case. """ matches = [] # Resolve scalar args exact match first if self.scalarpos: # Try resolve scalar arguments for formaltys in self.typemap: match_map = [] for i, (formal, actual) in enumerate(zip(formaltys, self.argtypes)): if actual is None: actual = np.asarray(self.args[i]).dtype match_map.append(actual == formal) if all(match_map): matches.append(formaltys) # No matching with exact match; try coercing the scalar arguments if not matches: matches = [] for formaltys in self.typemap: all_matches = all(actual is None or formal == actual for formal, actual in zip(formaltys, self.argtypes)) if all_matches: matches.append(formaltys) if not matches: raise TypeError("No matching version. GPU ufunc requires array " "arguments to have the exact types. This behaves " "like regular ufunc with casting='no'.") if len(matches) > 1: raise TypeError("Failed to resolve ufunc due to ambiguous " "signature. Too many untyped scalars. " "Use numpy dtype object to type tag.") # Try scalar arguments self.argtypes = matches[0] def _get_actual_args(self): """Return the actual arguments Casts scalar arguments to np.array. """ for i in self.scalarpos: self.arrays[i] = np.array([self.args[i]], dtype=self.argtypes[i]) return self.arrays def _broadcast(self, arys): """Perform numpy ufunc broadcasting """ shapelist = [a.shape for a in arys] shape = _multi_broadcast(*shapelist) for i, ary in enumerate(arys): if ary.shape == shape: pass else: if self.is_device_array(ary): arys[i] = self.broadcast_device(ary, shape) else: ax_differs = [ax for ax in range(len(shape)) if ax >= ary.ndim or ary.shape[ax] != shape[ax]] missingdim = len(shape) - len(ary.shape) strides = [0] * missingdim + list(ary.strides) for ax in ax_differs: strides[ax] = 0 strided = np.lib.stride_tricks.as_strided(ary, shape=shape, strides=strides) arys[i] = self.force_array_layout(strided) return arys def get_arguments(self): """Prepare and return the arguments for the ufunc. Does not call to_device(). """ self._fill_arrays() self._fill_argtypes() self._resolve_signature() arys = self._get_actual_args() return self._broadcast(arys) def get_function(self): """Returns (result_dtype, function) """ return self.typemap[self.argtypes] def is_device_array(self, obj): """Is the `obj` a device array? Override in subclass """ return False def as_device_array(self, obj): """Convert the `obj` to a device array Override in subclass Default implementation is an identity function """ return obj def broadcast_device(self, ary, shape): """Handles ondevice broadcasting Override in subclass to add support. """ raise NotImplementedError("broadcasting on device is not supported") def force_array_layout(self, ary): """Ensures array layout met device requirement. Override in sublcass """ return ary @classmethod def call(cls, typemap, args, kws): """Perform the entire ufunc call mechanism. """ # Handle keywords stream = kws.pop('stream', cls.DEFAULT_STREAM) out = kws.pop('out', None) if kws: warnings.warn("unrecognized keywords: %s" % ', '.join(kws)) # Begin call resolution cr = cls(typemap, args) args = cr.get_arguments() resty, func = cr.get_function() outshape = args[0].shape # Adjust output value if out is not None and cr.is_device_array(out): out = cr.as_device_array(out) def attempt_ravel(a): if cr.SUPPORT_DEVICE_SLICING: raise NotImplementedError try: # Call the `.ravel()` method return a.ravel() except NotImplementedError: # If it is not a device array if not cr.is_device_array(a): raise # For device array, retry ravel on the host by first # copying it back. else: hostary = cr.to_host(a, stream).ravel() return cr.to_device(hostary, stream) if args[0].ndim > 1: args = [attempt_ravel(a) for a in args] # Prepare argument on the device devarys = [] any_device = False for a in args: if cr.is_device_array(a): devarys.append(a) any_device = True else: dev_a = cr.to_device(a, stream=stream) devarys.append(dev_a) # Launch shape = args[0].shape if out is None: # No output is provided devout = cr.allocate_device_array(shape, resty, stream=stream) devarys.extend([devout]) cr.launch(func, shape[0], stream, devarys) if any_device: # If any of the arguments are on device, # Keep output on the device return devout.reshape(outshape) else: # Otherwise, transfer output back to host return devout.copy_to_host().reshape(outshape) elif cr.is_device_array(out): # If output is provided and it is a device array, # Return device array if out.ndim > 1: out = attempt_ravel(out) devout = out devarys.extend([devout]) cr.launch(func, shape[0], stream, devarys) return devout.reshape(outshape) else: # If output is provided and it is a host array, # Return host array assert out.shape == shape assert out.dtype == resty devout = cr.allocate_device_array(shape, resty, stream=stream) devarys.extend([devout]) cr.launch(func, shape[0], stream, devarys) return devout.copy_to_host(out, stream=stream).reshape(outshape) def to_device(self, hostary, stream): """Implement to device transfer Override in subclass """ raise NotImplementedError def to_host(self, devary, stream): """Implement to host transfer Override in subclass """ raise NotImplementedError def allocate_device_array(self, shape, dtype, stream): """Implements device allocation Override in subclass """ raise NotImplementedError def launch(self, func, count, stream, args): """Implements device function invocation Override in subclass """ raise NotImplementedError def to_dtype(ty): if isinstance(ty, types.EnumMember): ty = ty.dtype return np.dtype(str(ty)) class DeviceVectorize(_BaseUFuncBuilder): def __init__(self, func, identity=None, cache=False, targetoptions={}): if cache: raise TypeError("caching is not supported") for opt in targetoptions: if opt == 'nopython': warnings.warn("nopython kwarg for cuda target is redundant", RuntimeWarning) else: fmt = "Unrecognized options. " fmt += "cuda vectorize target does not support option: '%s'" raise KeyError(fmt % opt) self.py_func = func self.identity = parse_identity(identity) # { arg_dtype: (return_dtype), cudakernel } self.kernelmap = OrderedDict() @property def pyfunc(self): return self.py_func def add(self, sig=None): # compile core as device function args, return_type = sigutils.normalize_signature(sig) devfnsig = signature(return_type, *args) funcname = self.pyfunc.__name__ kernelsource = self._get_kernel_source(self._kernel_template, devfnsig, funcname) corefn, return_type = self._compile_core(devfnsig) glbl = self._get_globals(corefn) sig = signature(types.void, *([a[:] for a in args] + [return_type[:]])) exec(kernelsource, glbl) stager = glbl['__vectorized_%s' % funcname] kernel = self._compile_kernel(stager, sig) argdtypes = tuple(to_dtype(t) for t in devfnsig.args) resdtype = to_dtype(return_type) self.kernelmap[tuple(argdtypes)] = resdtype, kernel def build_ufunc(self): raise NotImplementedError def _get_kernel_source(self, template, sig, funcname): args = ['a%d' % i for i in range(len(sig.args))] fmts = dict(name=funcname, args=', '.join(args), argitems=', '.join('%s[__tid__]' % i for i in args)) return template.format(**fmts) def _compile_core(self, sig): raise NotImplementedError def _get_globals(self, corefn): raise NotImplementedError def _compile_kernel(self, fnobj, sig): raise NotImplementedError class DeviceGUFuncVectorize(_BaseUFuncBuilder): def __init__(self, func, sig, identity=None, cache=False, targetoptions={}, writable_args=()): if cache: raise TypeError("caching is not supported") if writable_args: raise TypeError("writable_args are not supported") # Allow nopython flag to be set. if not targetoptions.pop('nopython', True): raise TypeError("nopython flag must be True") # Are there any more target options? if targetoptions: opts = ', '.join([repr(k) for k in targetoptions.keys()]) fmt = "The following target options are not supported: {0}" raise TypeError(fmt.format(opts)) self.py_func = func self.identity = parse_identity(identity) self.signature = sig self.inputsig, self.outputsig = parse_signature(self.signature) # Maps from a tuple of input_dtypes to (output_dtypes, kernel) self.kernelmap = OrderedDict() @property def pyfunc(self): return self.py_func def add(self, sig=None): indims = [len(x) for x in self.inputsig] outdims = [len(x) for x in self.outputsig] args, return_type = sigutils.normalize_signature(sig) # It is only valid to specify types.none as a return type, or to not # specify the return type (where the "Python None" is the return type) valid_return_type = return_type in (types.none, None) if not valid_return_type: raise TypeError('guvectorized functions cannot return values: ' f'signature {sig} specifies {return_type} return ' 'type') funcname = self.py_func.__name__ src = expand_gufunc_template(self._kernel_template, indims, outdims, funcname, args) glbls = self._get_globals(sig) exec(src, glbls) fnobj = glbls['__gufunc_{name}'.format(name=funcname)] outertys = list(_determine_gufunc_outer_types(args, indims + outdims)) kernel = self._compile_kernel(fnobj, sig=tuple(outertys)) nout = len(outdims) dtypes = [np.dtype(str(t.dtype)) for t in outertys] indtypes = tuple(dtypes[:-nout]) outdtypes = tuple(dtypes[-nout:]) self.kernelmap[indtypes] = outdtypes, kernel def _compile_kernel(self, fnobj, sig): raise NotImplementedError def _get_globals(self, sig): raise NotImplementedError def _determine_gufunc_outer_types(argtys, dims): for at, nd in zip(argtys, dims): if isinstance(at, types.Array): yield at.copy(ndim=nd + 1) else: if nd > 0: raise ValueError("gufunc signature mismatch: ndim>0 for scalar") yield types.Array(dtype=at, ndim=1, layout='A') def expand_gufunc_template(template, indims, outdims, funcname, argtypes): """Expand gufunc source template """ argdims = indims + outdims argnames = ["arg{0}".format(i) for i in range(len(argdims))] checkedarg = "min({0})".format(', '.join(["{0}.shape[0]".format(a) for a in argnames])) inputs = [_gen_src_for_indexing(aref, adims, atype) for aref, adims, atype in zip(argnames, indims, argtypes)] outputs = [_gen_src_for_indexing(aref, adims, atype) for aref, adims, atype in zip(argnames[len(indims):], outdims, argtypes[len(indims):])] argitems = inputs + outputs src = template.format(name=funcname, args=', '.join(argnames), checkedarg=checkedarg, argitems=', '.join(argitems)) return src def _gen_src_for_indexing(aref, adims, atype): return "{aref}[{sliced}]".format(aref=aref, sliced=_gen_src_index(adims, atype)) def _gen_src_index(adims, atype): if adims > 0: return ','.join(['__tid__'] + [':'] * adims) elif isinstance(atype, types.Array) and atype.ndim - 1 == adims: # Special case for 0-nd in shape-signature but # 1d array in type signature. # Slice it so that the result has the same dimension. return '__tid__:(__tid__ + 1)' else: return '__tid__' class GUFuncEngine(object): '''Determine how to broadcast and execute a gufunc base on input shape and signature ''' @classmethod def from_signature(cls, signature): return cls(*parse_signature(signature)) def __init__(self, inputsig, outputsig): # signatures self.sin = inputsig self.sout = outputsig # argument count self.nin = len(self.sin) self.nout = len(self.sout) def schedule(self, ishapes): if len(ishapes) != self.nin: raise TypeError('invalid number of input argument') # associate symbol values for input signature symbolmap = {} outer_shapes = [] inner_shapes = [] for argn, (shape, symbols) in enumerate(zip(ishapes, self.sin)): argn += 1 # start from 1 for human inner_ndim = len(symbols) if len(shape) < inner_ndim: fmt = "arg #%d: insufficient inner dimension" raise ValueError(fmt % (argn,)) if inner_ndim: inner_shape = shape[-inner_ndim:] outer_shape = shape[:-inner_ndim] else: inner_shape = () outer_shape = shape for axis, (dim, sym) in enumerate(zip(inner_shape, symbols)): axis += len(outer_shape) if sym in symbolmap: if symbolmap[sym] != dim: fmt = "arg #%d: shape[%d] mismatch argument" raise ValueError(fmt % (argn, axis)) symbolmap[sym] = dim outer_shapes.append(outer_shape) inner_shapes.append(inner_shape) # solve output shape oshapes = [] for outsig in self.sout: oshape = [] for sym in outsig: oshape.append(symbolmap[sym]) oshapes.append(tuple(oshape)) # find the biggest outershape as looping dimension sizes = [reduce(operator.mul, s, 1) for s in outer_shapes] largest_i = np.argmax(sizes) loopdims = outer_shapes[largest_i] pinned = [False] * self.nin # same argument for each iteration for i, d in enumerate(outer_shapes): if d != loopdims: if d == (1,) or d == (): pinned[i] = True else: fmt = "arg #%d: outer dimension mismatch" raise ValueError(fmt % (i + 1,)) return GUFuncSchedule(self, inner_shapes, oshapes, loopdims, pinned) class GUFuncSchedule(object): def __init__(self, parent, ishapes, oshapes, loopdims, pinned): self.parent = parent # core shapes self.ishapes = ishapes self.oshapes = oshapes # looping dimension self.loopdims = loopdims self.loopn = reduce(operator.mul, loopdims, 1) # flags self.pinned = pinned self.output_shapes = [loopdims + s for s in oshapes] def __str__(self): import pprint attrs = 'ishapes', 'oshapes', 'loopdims', 'loopn', 'pinned' values = [(k, getattr(self, k)) for k in attrs] return pprint.pformat(dict(values)) class GeneralizedUFunc(object): def __init__(self, kernelmap, engine): self.kernelmap = kernelmap self.engine = engine self.max_blocksize = 2 ** 30 def __call__(self, *args, **kws): callsteps = self._call_steps(self.engine.nin, self.engine.nout, args, kws) indtypes, schedule, outdtypes, kernel = self._schedule( callsteps.inputs, callsteps.outputs) callsteps.adjust_input_types(indtypes) outputs = callsteps.prepare_outputs(schedule, outdtypes) inputs = callsteps.prepare_inputs() parameters = self._broadcast(schedule, inputs, outputs) callsteps.launch_kernel(kernel, schedule.loopn, parameters) return callsteps.post_process_outputs(outputs) def _schedule(self, inputs, outs): input_shapes = [a.shape for a in inputs] schedule = self.engine.schedule(input_shapes) # find kernel indtypes = tuple(i.dtype for i in inputs) try: outdtypes, kernel = self.kernelmap[indtypes] except KeyError: # No exact match, then use the first compatible. # This does not match the numpy dispatching exactly. # Later, we may just jit a new version for the missing signature. indtypes = self._search_matching_signature(indtypes) # Select kernel outdtypes, kernel = self.kernelmap[indtypes] # check output for sched_shape, out in zip(schedule.output_shapes, outs): if out is not None and sched_shape != out.shape: raise ValueError('output shape mismatch') return indtypes, schedule, outdtypes, kernel def _search_matching_signature(self, idtypes): """ Given the input types in `idtypes`, return a compatible sequence of types that is defined in `kernelmap`. Note: Ordering is guaranteed by `kernelmap` being a OrderedDict """ for sig in self.kernelmap.keys(): if all(np.can_cast(actual, desired) for actual, desired in zip(sig, idtypes)): return sig else: raise TypeError("no matching signature") def _broadcast(self, schedule, params, retvals): assert schedule.loopn > 0, "zero looping dimension" odim = 1 if not schedule.loopdims else schedule.loopn newparams = [] for p, cs in zip(params, schedule.ishapes): if not cs and p.size == 1: # Broadcast scalar input devary = self._broadcast_scalar_input(p, odim) newparams.append(devary) else: # Broadcast vector input newparams.append(self._broadcast_array(p, odim, cs)) newretvals = [] for retval, oshape in zip(retvals, schedule.oshapes): newretvals.append(retval.reshape(odim, *oshape)) return tuple(newparams) + tuple(newretvals) def _broadcast_array(self, ary, newdim, innerdim): newshape = (newdim,) + innerdim # No change in shape if ary.shape == newshape: return ary # Creating new dimension elif len(ary.shape) < len(newshape): assert newshape[-len(ary.shape):] == ary.shape, \ "cannot add dim and reshape at the same time" return self._broadcast_add_axis(ary, newshape) # Collapsing dimension else: return ary.reshape(*newshape) def _broadcast_add_axis(self, ary, newshape): raise NotImplementedError("cannot add new axis") def _broadcast_scalar_input(self, ary, shape): raise NotImplementedError class GUFuncCallSteps(metaclass=ABCMeta): """ Implements memory management and kernel launch operations for GUFunc calls. One instance of this class is instantiated for each call, and the instance is specific to the arguments given to the GUFunc call. The base class implements the overall logic; subclasses provide target-specific implementations of individual functions. """ # The base class uses these slots; subclasses may provide additional slots. __slots__ = [ 'outputs', 'inputs', '_copy_result_to_host', ] @abstractmethod def launch_kernel(self, kernel, nelem, args): """Implement the kernel launch""" @abstractmethod def is_device_array(self, obj): """ Return True if `obj` is a device array for this target, False otherwise. """ @abstractmethod def as_device_array(self, obj): """ Return `obj` as a device array on this target. May return `obj` directly if it is already on the target. """ @abstractmethod def to_device(self, hostary): """ Copy `hostary` to the device and return the device array. """ @abstractmethod def allocate_device_array(self, shape, dtype): """ Allocate a new uninitialized device array with the given shape and dtype. """ def __init__(self, nin, nout, args, kwargs): outputs = kwargs.get('out') # Ensure the user has passed a correct number of arguments if outputs is None and len(args) not in (nin, (nin + nout)): def pos_argn(n): return f'{n} positional argument{"s" * (n != 1)}' msg = (f'This gufunc accepts {pos_argn(nin)} (when providing ' f'input only) or {pos_argn(nin + nout)} (when providing ' f'input and output). Got {pos_argn(len(args))}.') raise TypeError(msg) if outputs is not None and len(args) > nin: raise ValueError("cannot specify argument 'out' as both positional " "and keyword") else: # If the user did not pass outputs either in the out kwarg or as # positional arguments, then we need to generate an initial list of # "placeholder" outputs using None as a sentry value outputs = [outputs] * nout # Ensure all output device arrays are Numba device arrays - for # example, any output passed in that supports the CUDA Array Interface # is converted to a Numba CUDA device array; others are left untouched. all_user_outputs_are_host = True self.outputs = [] for output in outputs: if self.is_device_array(output): self.outputs.append(self.as_device_array(output)) all_user_outputs_are_host = False else: self.outputs.append(output) all_host_arrays = not any([self.is_device_array(a) for a in args]) # - If any of the arguments are device arrays, we leave the output on # the device. self._copy_result_to_host = (all_host_arrays and all_user_outputs_are_host) # Normalize arguments - ensure they are either device- or host-side # arrays (as opposed to lists, tuples, etc). def normalize_arg(a): if self.is_device_array(a): convert = self.as_device_array else: convert = np.asarray return convert(a) normalized_args = [normalize_arg(a) for a in args] self.inputs = normalized_args[:nin] # Check if there are extra arguments for outputs. unused_inputs = normalized_args[nin:] if unused_inputs: self.outputs = unused_inputs def adjust_input_types(self, indtypes): """ Attempt to cast the inputs to the required types if necessary and if they are not device arrays. Side effect: Only affects the elements of `inputs` that require a type cast. """ for i, (ity, val) in enumerate(zip(indtypes, self.inputs)): if ity != val.dtype: if not hasattr(val, 'astype'): msg = ("compatible signature is possible by casting but " "{0} does not support .astype()").format(type(val)) raise TypeError(msg) # Cast types self.inputs[i] = val.astype(ity) def prepare_outputs(self, schedule, outdtypes): """ Returns a list of output parameters that all reside on the target device. Outputs that were passed-in to the GUFunc are used if they reside on the device; other outputs are allocated as necessary. """ outputs = [] for shape, dtype, output in zip(schedule.output_shapes, outdtypes, self.outputs): if output is None or self._copy_result_to_host: output = self.allocate_device_array(shape, dtype) outputs.append(output) return outputs def prepare_inputs(self): """ Returns a list of input parameters that all reside on the target device. """ def ensure_device(parameter): if self.is_device_array(parameter): convert = self.as_device_array else: convert = self.to_device return convert(parameter) return [ensure_device(p) for p in self.inputs] def post_process_outputs(self, outputs): """ Moves the given output(s) to the host if necessary. Returns a single value (e.g. an array) if there was one output, or a tuple of arrays if there were multiple. Although this feels a little jarring, it is consistent with the behavior of GUFuncs in general. """ if self._copy_result_to_host: outputs = [self.to_host(output, self_output) for output, self_output in zip(outputs, self.outputs)] elif self.outputs[0] is not None: outputs = self.outputs if len(outputs) == 1: return outputs[0] else: return tuple(outputs)