numba-cuda 0.0.0__py3-none-any.whl → 0.0.12__py3-none-any.whl

numba_cuda/numba/cuda/dispatcher.py

@@ -0,0 +1,1057 @@
+import numpy as np
+import os
+import sys
+import ctypes
+import functools
+
+from numba.core import config, serialize, sigutils, types, typing, utils
+from numba.core.caching import Cache, CacheImpl
+from numba.core.compiler_lock import global_compiler_lock
+from numba.core.dispatcher import Dispatcher
+from numba.core.errors import NumbaPerformanceWarning
+from numba.core.typing.typeof import Purpose, typeof
+
+from numba.cuda.api import get_current_device
+from numba.cuda.args import wrap_arg
+from numba.cuda.compiler import compile_cuda, CUDACompiler
+from numba.cuda.cudadrv import driver
+from numba.cuda.cudadrv.devices import get_context
+from numba.cuda.descriptor import cuda_target
+from numba.cuda.errors import (missing_launch_config_msg,
+                               normalize_kernel_dimensions)
+from numba.cuda import types as cuda_types
+
+from numba import cuda
+from numba import _dispatcher
+
+from warnings import warn
+
+cuda_fp16_math_funcs = ['hsin', 'hcos',
+                        'hlog', 'hlog10',
+                        'hlog2',
+                        'hexp', 'hexp10',
+                        'hexp2',
+                        'hsqrt', 'hrsqrt',
+                        'hfloor', 'hceil',
+                        'hrcp', 'hrint',
+                        'htrunc', 'hdiv']
+
+
+class _Kernel(serialize.ReduceMixin):
+    '''
+    CUDA Kernel specialized for a given set of argument types. When called, this
+    object launches the kernel on the device.
+    '''
+
+    @global_compiler_lock
+    def __init__(self, py_func, argtypes, link=None, debug=False,
+                 lineinfo=False, inline=False, fastmath=False, extensions=None,
+                 max_registers=None, opt=True, device=False):
+
+        if device:
+            raise RuntimeError('Cannot compile a device function as a kernel')
+
+        super().__init__()
+
+        # _DispatcherBase.nopython_signatures() expects this attribute to be
+        # present, because it assumes an overload is a CompileResult. In the
+        # CUDA target, _Kernel instances are stored instead, so we provide this
+        # attribute here to avoid duplicating nopython_signatures() in the CUDA
+        # target with slight modifications.
+        self.objectmode = False
+
+        # The finalizer constructed by _DispatcherBase._make_finalizer also
+        # expects overloads to be a CompileResult. It uses the entry_point to
+        # remove a CompileResult from a target context. However, since we never
+        # insert kernels into a target context (there is no need because they
+        # cannot be called by other functions, only through the dispatcher) it
+        # suffices to pretend we have an entry point of None.
+        self.entry_point = None
+
+        self.py_func = py_func
+        self.argtypes = argtypes
+        self.debug = debug
+        self.lineinfo = lineinfo
+        self.extensions = extensions or []
+
+        nvvm_options = {
+            'fastmath': fastmath,
+            'opt': 3 if opt else 0
+        }
+
+        cc = get_current_device().compute_capability
+        cres = compile_cuda(self.py_func, types.void, self.argtypes,
+                            debug=self.debug,
+                            lineinfo=lineinfo,
+                            inline=inline,
+                            fastmath=fastmath,
+                            nvvm_options=nvvm_options,
+                            cc=cc)
+        tgt_ctx = cres.target_context
+        code = self.py_func.__code__
+        filename = code.co_filename
+        linenum = code.co_firstlineno
+        lib, kernel = tgt_ctx.prepare_cuda_kernel(cres.library, cres.fndesc,
+                                                  debug, lineinfo, nvvm_options,
+                                                  filename, linenum,
+                                                  max_registers)
+
+        if not link:
+            link = []
+
+        # A kernel needs cooperative launch if grid_sync is being used.
+        self.cooperative = 'cudaCGGetIntrinsicHandle' in lib.get_asm_str()
+        # We need to link against cudadevrt if grid sync is being used.
+        if self.cooperative:
+            lib.needs_cudadevrt = True
+
+        res = [fn for fn in cuda_fp16_math_funcs
+               if (f'__numba_wrapper_{fn}' in lib.get_asm_str())]
+
+        if res:
+            # Path to the source containing the foreign function
+            basedir = os.path.dirname(os.path.abspath(__file__))
+            functions_cu_path = os.path.join(basedir,
+                                             'cpp_function_wrappers.cu')
+            link.append(functions_cu_path)
+
+        for filepath in link:
+            lib.add_linking_file(filepath)
+
+        # populate members
+        self.entry_name = kernel.name
+        self.signature = cres.signature
+        self._type_annotation = cres.type_annotation
+        self._codelibrary = lib
+        self.call_helper = cres.call_helper
+
+        # The following are referred to by the cache implementation. Note:
+        # - There are no referenced environments in CUDA.
+        # - Kernels don't have lifted code.
+        # - reload_init is only for parfors.
+        self.target_context = tgt_ctx
+        self.fndesc = cres.fndesc
+        self.environment = cres.environment
+        self._referenced_environments = []
+        self.lifted = []
+        self.reload_init = []
+
+    @property
+    def library(self):
+        return self._codelibrary
+
+    @property
+    def type_annotation(self):
+        return self._type_annotation
+
+    def _find_referenced_environments(self):
+        return self._referenced_environments
+
+    @property
+    def codegen(self):
+        return self.target_context.codegen()
+
+    @property
+    def argument_types(self):
+        return tuple(self.signature.args)
+
+    @classmethod
+    def _rebuild(cls, cooperative, name, signature, codelibrary,
+                 debug, lineinfo, call_helper, extensions):
+        """
+        Rebuild an instance.
+        """
+        instance = cls.__new__(cls)
+        # invoke parent constructor
+        super(cls, instance).__init__()
+        # populate members
+        instance.entry_point = None
+        instance.cooperative = cooperative
+        instance.entry_name = name
+        instance.signature = signature
+        instance._type_annotation = None
+        instance._codelibrary = codelibrary
+        instance.debug = debug
+        instance.lineinfo = lineinfo
+        instance.call_helper = call_helper
+        instance.extensions = extensions
+        return instance
+
+    def _reduce_states(self):
+        """
+        Reduce the instance for serialization.
+        Compiled definitions are serialized in PTX form.
+        Type annotation are discarded.
+        Thread, block and shared memory configuration are serialized.
+        Stream information is discarded.
+        """
+        return dict(cooperative=self.cooperative, name=self.entry_name,
+                    signature=self.signature, codelibrary=self._codelibrary,
+                    debug=self.debug, lineinfo=self.lineinfo,
+                    call_helper=self.call_helper, extensions=self.extensions)
+
+    def bind(self):
+        """
+        Force binding to current CUDA context
+        """
+        self._codelibrary.get_cufunc()
+
+    @property
+    def regs_per_thread(self):
+        '''
+        The number of registers used by each thread for this kernel.
+        '''
+        return self._codelibrary.get_cufunc().attrs.regs
+
+    @property
+    def const_mem_size(self):
+        '''
+        The amount of constant memory used by this kernel.
+        '''
+        return self._codelibrary.get_cufunc().attrs.const
+
+    @property
+    def shared_mem_per_block(self):
+        '''
+        The amount of shared memory used per block for this kernel.
+        '''
+        return self._codelibrary.get_cufunc().attrs.shared
+
+    @property
+    def max_threads_per_block(self):
+        '''
+        The maximum allowable threads per block.
+        '''
+        return self._codelibrary.get_cufunc().attrs.maxthreads
+
+    @property
+    def local_mem_per_thread(self):
+        '''
+        The amount of local memory used per thread for this kernel.
+        '''
+        return self._codelibrary.get_cufunc().attrs.local
+
+    def inspect_llvm(self):
+        '''
+        Returns the LLVM IR for this kernel.
+        '''
+        return self._codelibrary.get_llvm_str()
+
+    def inspect_asm(self, cc):
+        '''
+        Returns the PTX code for this kernel.
+        '''
+        return self._codelibrary.get_asm_str(cc=cc)
+
+    def inspect_sass_cfg(self):
+        '''
+        Returns the CFG of the SASS for this kernel.
+
+        Requires nvdisasm to be available on the PATH.
+        '''
+        return self._codelibrary.get_sass_cfg()
+
+    def inspect_sass(self):
+        '''
+        Returns the SASS code for this kernel.
+
+        Requires nvdisasm to be available on the PATH.
+        '''
+        return self._codelibrary.get_sass()
+
+    def inspect_types(self, file=None):
+        '''
+        Produce a dump of the Python source of this function annotated with the
+        corresponding Numba IR and type information. The dump is written to
+        *file*, or *sys.stdout* if *file* is *None*.
+        '''
+        if self._type_annotation is None:
+            raise ValueError("Type annotation is not available")
+
+        if file is None:
+            file = sys.stdout
+
+        print("%s %s" % (self.entry_name, self.argument_types), file=file)
+        print('-' * 80, file=file)
+        print(self._type_annotation, file=file)
+        print('=' * 80, file=file)
+
+    def max_cooperative_grid_blocks(self, blockdim, dynsmemsize=0):
+        '''
+        Calculates the maximum number of blocks that can be launched for this
+        kernel in a cooperative grid in the current context, for the given block
+        and dynamic shared memory sizes.
+
+        :param blockdim: Block dimensions, either as a scalar for a 1D block, or
+                         a tuple for 2D or 3D blocks.
+        :param dynsmemsize: Dynamic shared memory size in bytes.
+        :return: The maximum number of blocks in the grid.
+        '''
+        ctx = get_context()
+        cufunc = self._codelibrary.get_cufunc()
+
+        if isinstance(blockdim, tuple):
+            blockdim = functools.reduce(lambda x, y: x * y, blockdim)
+        active_per_sm = ctx.get_active_blocks_per_multiprocessor(cufunc,
+                                                                 blockdim,
+                                                                 dynsmemsize)
+        sm_count = ctx.device.MULTIPROCESSOR_COUNT
+        return active_per_sm * sm_count
+
+    def launch(self, args, griddim, blockdim, stream=0, sharedmem=0):
+        # Prepare kernel
+        cufunc = self._codelibrary.get_cufunc()
+
+        if self.debug:
+            excname = cufunc.name + "__errcode__"
+            excmem, excsz = cufunc.module.get_global_symbol(excname)
+            assert excsz == ctypes.sizeof(ctypes.c_int)
+            excval = ctypes.c_int()
+            excmem.memset(0, stream=stream)
+
+        # Prepare arguments
+        retr = []                       # hold functors for writeback
+
+        kernelargs = []
+        for t, v in zip(self.argument_types, args):
+            self._prepare_args(t, v, stream, retr, kernelargs)
+
+        if driver.USE_NV_BINDING:
+            zero_stream = driver.binding.CUstream(0)
+        else:
+            zero_stream = None
+
+        stream_handle = stream and stream.handle or zero_stream
+
+        # Invoke kernel
+        driver.launch_kernel(cufunc.handle,
+                             *griddim,
+                             *blockdim,
+                             sharedmem,
+                             stream_handle,
+                             kernelargs,
+                             cooperative=self.cooperative)
+
+        if self.debug:
+            driver.device_to_host(ctypes.addressof(excval), excmem, excsz)
+            if excval.value != 0:
+                # An error occurred
+                def load_symbol(name):
+                    mem, sz = cufunc.module.get_global_symbol("%s__%s__" %
+                                                              (cufunc.name,
+                                                               name))
+                    val = ctypes.c_int()
+                    driver.device_to_host(ctypes.addressof(val), mem, sz)
+                    return val.value
+
+                tid = [load_symbol("tid" + i) for i in 'zyx']
+                ctaid = [load_symbol("ctaid" + i) for i in 'zyx']
+                code = excval.value
+                exccls, exc_args, loc = self.call_helper.get_exception(code)
+                # Prefix the exception message with the source location
+                if loc is None:
+                    locinfo = ''
+                else:
+                    sym, filepath, lineno = loc
+                    filepath = os.path.abspath(filepath)
+                    locinfo = 'In function %r, file %s, line %s, ' % (sym,
+                                                                      filepath,
+                                                                      lineno,)
+                # Prefix the exception message with the thread position
+                prefix = "%stid=%s ctaid=%s" % (locinfo, tid, ctaid)
+                if exc_args:
+                    exc_args = ("%s: %s" % (prefix, exc_args[0]),) + \
+                        exc_args[1:]
+                else:
+                    exc_args = prefix,
+                raise exccls(*exc_args)
+
+        # retrieve auto converted arrays
+        for wb in retr:
+            wb()
+
+    def _prepare_args(self, ty, val, stream, retr, kernelargs):
+        """
+        Convert arguments to ctypes and append to kernelargs
+        """
+
+        # map the arguments using any extension you've registered
+        for extension in reversed(self.extensions):
+            ty, val = extension.prepare_args(
+                ty,
+                val,
+                stream=stream,
+                retr=retr)
+
+        if isinstance(ty, types.Array):
+            devary = wrap_arg(val).to_device(retr, stream)
+
+            c_intp = ctypes.c_ssize_t
+
+            meminfo = ctypes.c_void_p(0)
+            parent = ctypes.c_void_p(0)
+            nitems = c_intp(devary.size)
+            itemsize = c_intp(devary.dtype.itemsize)
+
+            ptr = driver.device_pointer(devary)
+
+            if driver.USE_NV_BINDING:
+                ptr = int(ptr)
+
+            data = ctypes.c_void_p(ptr)
+
+            kernelargs.append(meminfo)
+            kernelargs.append(parent)
+            kernelargs.append(nitems)
+            kernelargs.append(itemsize)
+            kernelargs.append(data)
+            for ax in range(devary.ndim):
+                kernelargs.append(c_intp(devary.shape[ax]))
+            for ax in range(devary.ndim):
+                kernelargs.append(c_intp(devary.strides[ax]))
+
+        elif isinstance(ty, types.Integer):
+            cval = getattr(ctypes, "c_%s" % ty)(val)
+            kernelargs.append(cval)
+
+        elif ty == types.float16:
+            cval = ctypes.c_uint16(np.float16(val).view(np.uint16))
+            kernelargs.append(cval)
+
+        elif ty == types.float64:
+            cval = ctypes.c_double(val)
+            kernelargs.append(cval)
+
+        elif ty == types.float32:
+            cval = ctypes.c_float(val)
+            kernelargs.append(cval)
+
+        elif ty == types.boolean:
+            cval = ctypes.c_uint8(int(val))
+            kernelargs.append(cval)
+
+        elif ty == types.complex64:
+            kernelargs.append(ctypes.c_float(val.real))
+            kernelargs.append(ctypes.c_float(val.imag))
+
+        elif ty == types.complex128:
+            kernelargs.append(ctypes.c_double(val.real))
+            kernelargs.append(ctypes.c_double(val.imag))
+
+        elif isinstance(ty, (types.NPDatetime, types.NPTimedelta)):
+            kernelargs.append(ctypes.c_int64(val.view(np.int64)))
+
+        elif isinstance(ty, types.Record):
+            devrec = wrap_arg(val).to_device(retr, stream)
+            ptr = devrec.device_ctypes_pointer
+            if driver.USE_NV_BINDING:
+                ptr = ctypes.c_void_p(int(ptr))
+            kernelargs.append(ptr)
+
+        elif isinstance(ty, types.BaseTuple):
+            assert len(ty) == len(val)
+            for t, v in zip(ty, val):
+                self._prepare_args(t, v, stream, retr, kernelargs)
+
+        elif isinstance(ty, types.EnumMember):
+            try:
+                self._prepare_args(
+                    ty.dtype, val.value, stream, retr, kernelargs
+                )
+            except NotImplementedError:
+                raise NotImplementedError(ty, val)
+
+        else:
+            raise NotImplementedError(ty, val)
+
+
+class ForAll(object):
+    def __init__(self, dispatcher, ntasks, tpb, stream, sharedmem):
+        if ntasks < 0:
+            raise ValueError("Can't create ForAll with negative task count: %s"
+                             % ntasks)
+        self.dispatcher = dispatcher
+        self.ntasks = ntasks
+        self.thread_per_block = tpb
+        self.stream = stream
+        self.sharedmem = sharedmem
+
+    def __call__(self, *args):
+        if self.ntasks == 0:
+            return
+
+        if self.dispatcher.specialized:
+            specialized = self.dispatcher
+        else:
+            specialized = self.dispatcher.specialize(*args)
+        blockdim = self._compute_thread_per_block(specialized)
+        griddim = (self.ntasks + blockdim - 1) // blockdim
+
+        return specialized[griddim, blockdim, self.stream,
+                           self.sharedmem](*args)
+
+    def _compute_thread_per_block(self, dispatcher):
+        tpb = self.thread_per_block
+        # Prefer user-specified config
+        if tpb != 0:
+            return tpb
+        # Else, ask the driver to give a good config
+        else:
+            ctx = get_context()
+            # Dispatcher is specialized, so there's only one definition - get
+            # it so we can get the cufunc from the code library
+            kernel = next(iter(dispatcher.overloads.values()))
+            kwargs = dict(
+                func=kernel._codelibrary.get_cufunc(),
+                b2d_func=0,     # dynamic-shared memory is constant to blksz
+                memsize=self.sharedmem,
+                blocksizelimit=1024,
+            )
+            _, tpb = ctx.get_max_potential_block_size(**kwargs)
+            return tpb
+
+
+class _LaunchConfiguration:
+    def __init__(self, dispatcher, griddim, blockdim, stream, sharedmem):
+        self.dispatcher = dispatcher
+        self.griddim = griddim
+        self.blockdim = blockdim
+        self.stream = stream
+        self.sharedmem = sharedmem
+
+        if config.CUDA_LOW_OCCUPANCY_WARNINGS:
+            # Warn when the grid has fewer than 128 blocks. This number is
+            # chosen somewhat heuristically - ideally the minimum is 2 times
+            # the number of SMs, but the number of SMs varies between devices -
+            # some very small GPUs might only have 4 SMs, but an H100-SXM5 has
+            # 132. In general kernels should be launched with large grids
+            # (hundreds or thousands of blocks), so warning when fewer than 128
+            # blocks are used will likely catch most beginner errors, where the
+            # grid tends to be very small (single-digit or low tens of blocks).
+            min_grid_size = 128
+            grid_size = griddim[0] * griddim[1] * griddim[2]
+            if grid_size < min_grid_size:
+                msg = (f"Grid size {grid_size} will likely result in GPU "
+                       "under-utilization due to low occupancy.")
+                warn(NumbaPerformanceWarning(msg))
+
+    def __call__(self, *args):
+        return self.dispatcher.call(args, self.griddim, self.blockdim,
+                                    self.stream, self.sharedmem)
+
+
+class CUDACacheImpl(CacheImpl):
+    def reduce(self, kernel):
+        return kernel._reduce_states()
+
+    def rebuild(self, target_context, payload):
+        return _Kernel._rebuild(**payload)
+
+    def check_cachable(self, cres):
+        # CUDA Kernels are always cachable - the reasons for an entity not to
+        # be cachable are:
+        #
+        # - The presence of lifted loops, or
+        # - The presence of dynamic globals.
+        #
+        # neither of which apply to CUDA kernels.
+        return True
+
+
+class CUDACache(Cache):
+    """
+    Implements a cache that saves and loads CUDA kernels and compile results.
+    """
+    _impl_class = CUDACacheImpl
+
+    def load_overload(self, sig, target_context):
+        # Loading an overload refreshes the context to ensure it is
+        # initialized. To initialize the correct (i.e. CUDA) target, we need to
+        # enforce that the current target is the CUDA target.
+        from numba.core.target_extension import target_override
+        with target_override('cuda'):
+            return super().load_overload(sig, target_context)
+
+
+class CUDADispatcher(Dispatcher, serialize.ReduceMixin):
+    '''
+    CUDA Dispatcher object. When configured and called, the dispatcher will
+    specialize itself for the given arguments (if no suitable specialized
+    version already exists) & compute capability, and launch on the device
+    associated with the current context.
+
+    Dispatcher objects are not to be constructed by the user, but instead are
+    created using the :func:`numba.cuda.jit` decorator.
+    '''
+
+    # Whether to fold named arguments and default values. Default values are
+    # presently unsupported on CUDA, so we can leave this as False in all
+    # cases.
+    _fold_args = False
+
+    targetdescr = cuda_target
+
+    def __init__(self, py_func, targetoptions, pipeline_class=CUDACompiler):
+        super().__init__(py_func, targetoptions=targetoptions,
+                         pipeline_class=pipeline_class)
+
+        # The following properties are for specialization of CUDADispatchers. A
+        # specialized CUDADispatcher is one that is compiled for exactly one
+        # set of argument types, and bypasses some argument type checking for
+        # faster kernel launches.
+
+        # Is this a specialized dispatcher?
+        self._specialized = False
+
+        # If we produced specialized dispatchers, we cache them for each set of
+        # argument types
+        self.specializations = {}
+
+    @property
+    def _numba_type_(self):
+        return cuda_types.CUDADispatcher(self)
+
+    def enable_caching(self):
+        self._cache = CUDACache(self.py_func)
+
+    @functools.lru_cache(maxsize=128)
+    def configure(self, griddim, blockdim, stream=0, sharedmem=0):
+        griddim, blockdim = normalize_kernel_dimensions(griddim, blockdim)
+        return _LaunchConfiguration(self, griddim, blockdim, stream, sharedmem)
+
+    def __getitem__(self, args):
+        if len(args) not in [2, 3, 4]:
+            raise ValueError('must specify at least the griddim and blockdim')
+        return self.configure(*args)
+
+    def forall(self, ntasks, tpb=0, stream=0, sharedmem=0):
+        """Returns a 1D-configured dispatcher for a given number of tasks.
+
+        This assumes that:
+
+        - the kernel maps the Global Thread ID ``cuda.grid(1)`` to tasks on a
+          1-1 basis.
+        - the kernel checks that the Global Thread ID is upper-bounded by
+          ``ntasks``, and does nothing if it is not.
+
+        :param ntasks: The number of tasks.
+        :param tpb: The size of a block. An appropriate value is chosen if this
+                    parameter is not supplied.
+        :param stream: The stream on which the configured dispatcher will be
+                       launched.
+        :param sharedmem: The number of bytes of dynamic shared memory required
+                          by the kernel.
+        :return: A configured dispatcher, ready to launch on a set of
+                 arguments."""
+
+        return ForAll(self, ntasks, tpb=tpb, stream=stream, sharedmem=sharedmem)
+
+    @property
+    def extensions(self):
+        '''
+        A list of objects that must have a `prepare_args` function. When a
+        specialized kernel is called, each argument will be passed through
+        to the `prepare_args` (from the last object in this list to the
+        first). The arguments to `prepare_args` are:
+
+        - `ty` the numba type of the argument
+        - `val` the argument value itself
+        - `stream` the CUDA stream used for the current call to the kernel
+        - `retr` a list of zero-arg functions that you may want to append
+          post-call cleanup work to.
+
+        The `prepare_args` function must return a tuple `(ty, val)`, which
+        will be passed in turn to the next right-most `extension`. After all
+        the extensions have been called, the resulting `(ty, val)` will be
+        passed into Numba's default argument marshalling logic.
+        '''
+        return self.targetoptions.get('extensions')
+
+    def __call__(self, *args, **kwargs):
+        # An attempt to launch an unconfigured kernel
+        raise ValueError(missing_launch_config_msg)
+
+    def call(self, args, griddim, blockdim, stream, sharedmem):
+        '''
+        Compile if necessary and invoke this kernel with *args*.
+        '''
+        if self.specialized:
+            kernel = next(iter(self.overloads.values()))
+        else:
+            kernel = _dispatcher.Dispatcher._cuda_call(self, *args)
+
+        kernel.launch(args, griddim, blockdim, stream, sharedmem)
+
+    def _compile_for_args(self, *args, **kws):
+        # Based on _DispatcherBase._compile_for_args.
+        assert not kws
+        argtypes = [self.typeof_pyval(a) for a in args]
+        return self.compile(tuple(argtypes))
+
+    def typeof_pyval(self, val):
+        # Based on _DispatcherBase.typeof_pyval, but differs from it to support
+        # the CUDA Array Interface.
+        try:
+            return typeof(val, Purpose.argument)
+        except ValueError:
+            if cuda.is_cuda_array(val):
+                # When typing, we don't need to synchronize on the array's
+                # stream - this is done when the kernel is launched.
+                return typeof(cuda.as_cuda_array(val, sync=False),
+                              Purpose.argument)
+            else:
+                raise
+
+    def specialize(self, *args):
+        '''
+        Create a new instance of this dispatcher specialized for the given
+        *args*.
+        '''
+        cc = get_current_device().compute_capability
+        argtypes = tuple(
+            [self.typingctx.resolve_argument_type(a) for a in args])
+        if self.specialized:
+            raise RuntimeError('Dispatcher already specialized')
+
+        specialization = self.specializations.get((cc, argtypes))
+        if specialization:
+            return specialization
+
+        targetoptions = self.targetoptions
+        specialization = CUDADispatcher(self.py_func,
+                                        targetoptions=targetoptions)
+        specialization.compile(argtypes)
+        specialization.disable_compile()
+        specialization._specialized = True
+        self.specializations[cc, argtypes] = specialization
+        return specialization
+
+    @property
+    def specialized(self):
+        """
+        True if the Dispatcher has been specialized.
+        """
+        return self._specialized
+
+    def get_regs_per_thread(self, signature=None):
+        '''
+        Returns the number of registers used by each thread in this kernel for
+        the device in the current context.
+
+        :param signature: The signature of the compiled kernel to get register
+                          usage for. This may be omitted for a specialized
+                          kernel.
+        :return: The number of registers used by the compiled variant of the
+                 kernel for the given signature and current device.
+        '''
+        if signature is not None:
+            return self.overloads[signature.args].regs_per_thread
+        if self.specialized:
+            return next(iter(self.overloads.values())).regs_per_thread
+        else:
+            return {sig: overload.regs_per_thread
+                    for sig, overload in self.overloads.items()}
+
+    def get_const_mem_size(self, signature=None):
+        '''
+        Returns the size in bytes of constant memory used by this kernel for
+        the device in the current context.
+
+        :param signature: The signature of the compiled kernel to get constant
+                          memory usage for. This may be omitted for a
+                          specialized kernel.
+        :return: The size in bytes of constant memory allocated by the
+                 compiled variant of the kernel for the given signature and
+                 current device.
+        '''
+        if signature is not None:
+            return self.overloads[signature.args].const_mem_size
+        if self.specialized:
+            return next(iter(self.overloads.values())).const_mem_size
+        else:
+            return {sig: overload.const_mem_size
+                    for sig, overload in self.overloads.items()}
+
+    def get_shared_mem_per_block(self, signature=None):
+        '''
+        Returns the size in bytes of statically allocated shared memory
+        for this kernel.
+
+        :param signature: The signature of the compiled kernel to get shared
+                          memory usage for. This may be omitted for a
+                          specialized kernel.
+        :return: The amount of shared memory allocated by the compiled variant
+                 of the kernel for the given signature and current device.
+        '''
+        if signature is not None:
+            return self.overloads[signature.args].shared_mem_per_block
+        if self.specialized:
+            return next(iter(self.overloads.values())).shared_mem_per_block
+        else:
+            return {sig: overload.shared_mem_per_block
+                    for sig, overload in self.overloads.items()}
+
+    def get_max_threads_per_block(self, signature=None):
+        '''
+        Returns the maximum allowable number of threads per block
+        for this kernel. Exceeding this threshold will result in
+        the kernel failing to launch.
+
+        :param signature: The signature of the compiled kernel to get the max
+                          threads per block for. This may be omitted for a
+                          specialized kernel.
+        :return: The maximum allowable threads per block for the compiled
+                 variant of the kernel for the given signature and current
+                 device.
+        '''
+        if signature is not None:
+            return self.overloads[signature.args].max_threads_per_block
+        if self.specialized:
+            return next(iter(self.overloads.values())).max_threads_per_block
+        else:
+            return {sig: overload.max_threads_per_block
+                    for sig, overload in self.overloads.items()}
+
+    def get_local_mem_per_thread(self, signature=None):
+        '''
+        Returns the size in bytes of local memory per thread
+        for this kernel.
+
+        :param signature: The signature of the compiled kernel to get local
+                          memory usage for. This may be omitted for a
+                          specialized kernel.
+        :return: The amount of local memory allocated by the compiled variant
+                 of the kernel for the given signature and current device.
+        '''
+        if signature is not None:
+            return self.overloads[signature.args].local_mem_per_thread
+        if self.specialized:
+            return next(iter(self.overloads.values())).local_mem_per_thread
+        else:
+            return {sig: overload.local_mem_per_thread
+                    for sig, overload in self.overloads.items()}
+
+    def get_call_template(self, args, kws):
+        # Originally copied from _DispatcherBase.get_call_template. This
+        # version deviates slightly from the _DispatcherBase version in order
+        # to force casts when calling device functions. See e.g.
+        # TestDeviceFunc.test_device_casting, added in PR #7496.
+        """
+        Get a typing.ConcreteTemplate for this dispatcher and the given
+        *args* and *kws* types.  This allows resolution of the return type.
+
+        A (template, pysig, args, kws) tuple is returned.
+        """
+        # Ensure an exactly-matching overload is available if we can
+        # compile. We proceed with the typing even if we can't compile
+        # because we may be able to force a cast on the caller side.
+        if self._can_compile:
+            self.compile_device(tuple(args))
+
+        # Create function type for typing
+        func_name = self.py_func.__name__
+        name = "CallTemplate({0})".format(func_name)
+
+        call_template = typing.make_concrete_template(
+            name, key=func_name, signatures=self.nopython_signatures)
+        pysig = utils.pysignature(self.py_func)
+
+        return call_template, pysig, args, kws
+
+    def compile_device(self, args, return_type=None):
+        """Compile the device function for the given argument types.
+
+        Each signature is compiled once by caching the compiled function inside
+        this object.
+
+        Returns the `CompileResult`.
+        """
+        if args not in self.overloads:
+            with self._compiling_counter:
+
+                debug = self.targetoptions.get('debug')
+                lineinfo = self.targetoptions.get('lineinfo')
+                inline = self.targetoptions.get('inline')
+                fastmath = self.targetoptions.get('fastmath')
+
+                nvvm_options = {
+                    'opt': 3 if self.targetoptions.get('opt') else 0,
+                    'fastmath': fastmath
+                }
+
+                cc = get_current_device().compute_capability
+                cres = compile_cuda(self.py_func, return_type, args,
+                                    debug=debug,
+                                    lineinfo=lineinfo,
+                                    inline=inline,
+                                    fastmath=fastmath,
+                                    nvvm_options=nvvm_options,
+                                    cc=cc)
+                self.overloads[args] = cres
+
+                cres.target_context.insert_user_function(cres.entry_point,
+                                                         cres.fndesc,
+                                                         [cres.library])
+        else:
+            cres = self.overloads[args]
+
+        return cres
+
+    def add_overload(self, kernel, argtypes):
+        c_sig = [a._code for a in argtypes]
+        self._insert(c_sig, kernel, cuda=True)
+        self.overloads[argtypes] = kernel
+
+    def compile(self, sig):
+        '''
+        Compile and bind to the current context a version of this kernel
+        specialized for the given signature.
+        '''
+        argtypes, return_type = sigutils.normalize_signature(sig)
+        assert return_type is None or return_type == types.none
+
+        # Do we already have an in-memory compiled kernel?
+        if self.specialized:
+            return next(iter(self.overloads.values()))
+        else:
+            kernel = self.overloads.get(argtypes)
+            if kernel is not None:
+                return kernel
+
+        # Can we load from the disk cache?
+        kernel = self._cache.load_overload(sig, self.targetctx)
+
+        if kernel is not None:
+            self._cache_hits[sig] += 1
+        else:
+            # We need to compile a new kernel
+            self._cache_misses[sig] += 1
+            if not self._can_compile:
+                raise RuntimeError("Compilation disabled")
+
+            kernel = _Kernel(self.py_func, argtypes, **self.targetoptions)
+            # We call bind to force codegen, so that there is a cubin to cache
+            kernel.bind()
+            self._cache.save_overload(sig, kernel)
+
+        self.add_overload(kernel, argtypes)
+
+        return kernel
+
+    def inspect_llvm(self, signature=None):
+        '''
+        Return the LLVM IR for this kernel.
+
+        :param signature: A tuple of argument types.
+        :return: The LLVM IR for the given signature, or a dict of LLVM IR
+                 for all previously-encountered signatures.
+
+        '''
+        device = self.targetoptions.get('device')
+        if signature is not None:
+            if device:
+                return self.overloads[signature].library.get_llvm_str()
+            else:
+                return self.overloads[signature].inspect_llvm()
+        else:
+            if device:
+                return {sig: overload.library.get_llvm_str()
+                        for sig, overload in self.overloads.items()}
+            else:
+                return {sig: overload.inspect_llvm()
+                        for sig, overload in self.overloads.items()}
+
+    def inspect_asm(self, signature=None):
+        '''
+        Return this kernel's PTX assembly code for for the device in the
+        current context.
+
+        :param signature: A tuple of argument types.
+        :return: The PTX code for the given signature, or a dict of PTX codes
+                 for all previously-encountered signatures.
+        '''
+        cc = get_current_device().compute_capability
+        device = self.targetoptions.get('device')
+        if signature is not None:
+            if device:
+                return self.overloads[signature].library.get_asm_str(cc)
+            else:
+                return self.overloads[signature].inspect_asm(cc)
+        else:
+            if device:
+                return {sig: overload.library.get_asm_str(cc)
+                        for sig, overload in self.overloads.items()}
+            else:
+                return {sig: overload.inspect_asm(cc)
+                        for sig, overload in self.overloads.items()}
+
+    def inspect_sass_cfg(self, signature=None):
+        '''
+        Return this kernel's CFG for the device in the current context.
+
+        :param signature: A tuple of argument types.
+        :return: The CFG for the given signature, or a dict of CFGs
+                 for all previously-encountered signatures.
+
+        The CFG for the device in the current context is returned.
+
+        Requires nvdisasm to be available on the PATH.
+        '''
+        if self.targetoptions.get('device'):
+            raise RuntimeError('Cannot get the CFG of a device function')
+
+        if signature is not None:
+            return self.overloads[signature].inspect_sass_cfg()
+        else:
+            return {sig: defn.inspect_sass_cfg()
+                    for sig, defn in self.overloads.items()}
+
+    def inspect_sass(self, signature=None):
+        '''
+        Return this kernel's SASS assembly code for for the device in the
+        current context.
+
+        :param signature: A tuple of argument types.
+        :return: The SASS code for the given signature, or a dict of SASS codes
+                 for all previously-encountered signatures.
+
+        SASS for the device in the current context is returned.
+
+        Requires nvdisasm to be available on the PATH.
+        '''
+        if self.targetoptions.get('device'):
+            raise RuntimeError('Cannot inspect SASS of a device function')
+
+        if signature is not None:
+            return self.overloads[signature].inspect_sass()
+        else:
+            return {sig: defn.inspect_sass()
+                    for sig, defn in self.overloads.items()}
+
+    def inspect_types(self, file=None):
+        '''
+        Produce a dump of the Python source of this function annotated with the
+        corresponding Numba IR and type information. The dump is written to
+        *file*, or *sys.stdout* if *file* is *None*.
+        '''
+        if file is None:
+            file = sys.stdout
+
+        for _, defn in self.overloads.items():
+            defn.inspect_types(file=file)
+
+    @classmethod
+    def _rebuild(cls, py_func, targetoptions):
+        """
+        Rebuild an instance.
+        """
+        instance = cls(py_func, targetoptions)
+        return instance
+
+    def _reduce_states(self):
+        """
+        Reduce the instance for serialization.
+        Compiled definitions are discarded.
+        """
+        return dict(py_func=self.py_func,
+                    targetoptions=self.targetoptions)