deviceufunc.py

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