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*. ''' if self.specialized: raise RuntimeError('Dispatcher already specialized') cc = get_current_device().compute_capability argtypes = tuple(self.typeof_pyval(a) for a in args) 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)