numba-cuda 0.22.0__cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/deviceufunc.py

@@ -0,0 +1,1021 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+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 tokenize
+import string
+
+import numpy as np
+
+from numba.cuda.np.ufunc.ufuncbuilder import _BaseUFuncBuilder, parse_identity
+from numba.cuda import types
+from numba.cuda.typing import signature
+from numba.cuda.core import sigutils
+
+
+def parse_signature(sig):
+    """Parse generalized ufunc signature.
+
+    NOTE: ',' (COMMA) is a delimiter; not separator.
+          This means trailing comma is legal.
+    """
+
+    def stripws(s):
+        return "".join(c for c in s if c not in string.whitespace)
+
+    def tokenizer(src):
+        def readline():
+            yield src
+
+        gen = readline()
+        return tokenize.generate_tokens(lambda: next(gen))
+
+    def parse(src):
+        tokgen = tokenizer(src)
+        while True:
+            tok = next(tokgen)
+            if tok[1] == "(":
+                symbols = []
+                while True:
+                    tok = next(tokgen)
+                    if tok[1] == ")":
+                        break
+                    elif tok[0] == tokenize.NAME:
+                        symbols.append(tok[1])
+                    elif tok[1] == ",":
+                        continue
+                    else:
+                        raise ValueError('bad token in signature "%s"' % tok[1])
+                yield tuple(symbols)
+                tok = next(tokgen)
+                if tok[1] == ",":
+                    continue
+                elif tokenize.ISEOF(tok[0]):
+                    break
+            elif tokenize.ISEOF(tok[0]):
+                break
+            else:
+                raise ValueError('bad token in signature "%s"' % tok[1])
+
+    ins, _, outs = stripws(sig).partition("->")
+    inputs = list(parse(ins))
+    outputs = list(parse(outs))
+
+    # check that all output symbols are defined in the inputs
+    isym = set()
+    osym = set()
+    for grp in inputs:
+        isym |= set(grp)
+    for grp in outputs:
+        osym |= set(grp)
+
+    diff = osym.difference(isym)
+    if diff:
+        raise NameError("undefined output symbols: %s" % ",".join(sorted(diff)))
+
+    return inputs, outputs
+
+
+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)