numba_vulkan.dispatcher#

Dispatcher and jit decorator of the Vulkan target.

Classes#

VulkanDispatcher

A Python function compiled lazily for Vulkan.

Functions#

`raise_kernel_error`(name, code[, asynchronous])	Raise the exception that a kernel reported.
`jit`([pyfunc, fastmath, narrow_math, boundscheck, ...])	Compile a Python function for Vulkan.

Module Contents#

class numba_vulkan.dispatcher.VulkanDispatcher(py_func, targetoptions=None)#

A Python function compiled lazily for Vulkan.

It is launched as a kernel through forall, and it can be called from other Vulkan-compiled functions, in which case it is inlined. Instances are normally created with jit.

Parameters:

py_func (function) – The Python function.
targetoptions (dict, optional) – Options; fastmath, narrow_math and boundscheck are used.

Variables:

py_func (function) – The Python function.
fastmath (bool) – Whether speed is preferred over accuracy; see jit.
narrow_math (bool) – Whether float64 transcendental functions are computed in float32 when libclc is unavailable.
boundscheck (bool) – Whether array indices are checked.
narrow (bool or None) – Whether 64-bit types are narrowed to 32 bits: always, never, or (None) on devices that lack them; see jit.
overloads (dict) – Compile results by tuple of argument types.

Examples

>>> import numpy as np
>>> import numba_vulkan as nv
>>> @nv.jit
... def double(x, out):
...     i = nv.global_id(0)
...     if i < x.shape[0]:
...         out[i] = 2 * x[i]
>>> x = np.arange(4, dtype=np.float32)
>>> out = np.zeros_like(x)
>>> double.forall(x.size)(x, out)
>>> out
array([0., 2., 4., 6.], dtype=float32)

targetdescr#

property overloads#

Compile results by tuple of argument types.

Functions compiled for different device features (see numba_vulkan.narrowing.Mode) are kept apart; the property gives those of the mode being compiled in.

Return type:: dict

py_func#

targetoptions#

narrow_math#

fastmath#

boundscheck#

error_model#

narrow#

cache#

property is_compiling#

Whether a compilation of this function is in progress.

Numba’s type inference uses this to detect recursion.

Return type:: bool

property nopython_signatures#

Signatures of all specialisations compiled so far.

Return type:: list of numba.core.typing.Signature

compile_device(args, return_type=None)#

Compile for args as a function callable from other kernels.

Parameters:

args (tuple of numba.types.Type) – Argument types.
return_type (numba.types.Type, optional) – Return type; inferred when omitted.

Return type:

VulkanCompileResult

fold_argument_types(args, kws)#

Turn keyword, default and star arguments into positional types.

Parameters:

args (sequence of numba.types.Type) – Positional argument types of a call.
kws (dict) – Keyword argument types of the call.

Returns:

pysig (inspect.Signature) – Signature of the Python function.
args (tuple of numba.types.Type) – One type per parameter of the Python function.

get_call_template(args, kws)#

Provide the typing template for a call from compiled code.

Compiles the specialisation for the argument types if necessary.

Parameters:

args (sequence of numba.types.Type) – Positional argument types.
kws (dict) – Keyword argument types.

Returns:

template (type) – A concrete template holding the compiled signatures.
pysig (inspect.Signature) – Signature of the Python function.
args (tuple of numba.types.Type) – Folded argument types.
kws (dict) – Always empty; keywords were folded into args.

get_overload(sig)#

Key identifying the compiled specialisation for a signature.

Parameters:: sig (numba.core.typing.Signature or tuple of numba.types.Type) – A signature or its argument types.
Returns:: The entry point of the VulkanCompileResult.
Return type:: int
Raises:: KeyError – If the specialisation has not been compiled.

get_compile_result(sig)#

Compile for a signature and return the result.

Parameters:: sig (numba.core.typing.Signature) – The signature.
Return type:: VulkanCompileResult

compile(argtypes, ndim=1, mode=narrowing.Mode(), local_size=None, aliases=())#

Compile (or fetch) the kernel specialisation for argtypes.

Parameters:

argtypes (tuple of numba.types.Type) – Argument types; arrays are bound to consecutive storage buffers.
ndim (int) – Dimensionality of the dispatch grid.
mode (numba_vulkan.narrowing.Mode) – The 64-bit types the kernel must do without.
local_size (tuple of int, optional) – Workgroup size; by default 64 invocations, shaped by ndim.
aliases (tuple of tuple) – (index, earlier) for array arguments that share the buffer of an earlier one; see numba_vulkan.compiler.compile_kernel.

Return type:

CompiledKernel

forall(extent, device=None, local_size=None, stream=None)#

Bind a dispatch grid, like numba.cuda’s kernel.forall.

Parameters:

extent (int or tuple of int) – Number of invocations along each of up to three axes. It is rounded up to whole workgroups, so kernels must bounds-check numba_vulkan.stubs.global_id().
device (int, str or None) – Device to run on; defaults to the selected device.
local_size (int or tuple of int, optional) – Workgroup size. By default 64 invocations: (64,), (8, 8) or (4, 4, 4) depending on the grid.
stream (numba_vulkan.runtime.Stream, optional) – Enqueue the launch on a stream; see numba_vulkan.runtime.Stream. NumPy arrays are copied on the stream and hold the results after it is synchronized.

Returns:

Call it with the kernel arguments to run.

Return type:

callable

numba_vulkan.dispatcher.raise_kernel_error(name, code, asynchronous=False)#

Raise the exception that a kernel reported.

Parameters:

name (str) – Name of the kernel.
code (int) – Status code left by the kernel.
asynchronous (bool) – Whether the kernel was launched without waiting for it, so that the exception is raised later, by a synchronisation.

Raises:

Exception – The exception registered for code, with a note saying where in the kernel it was raised.

numba_vulkan.dispatcher.jit(pyfunc=None, *, fastmath=False, narrow_math=False, boundscheck=False, error_model='numpy', narrow=None, cache=False, **options)#

Compile a Python function for Vulkan.

Parameters:

pyfunc (function, optional) – The function. Kernels must return None and write into arrays; functions called from kernels may return values.
fastmath (bool) – Trade accuracy for speed: float32 math functions use the device’s built-in versions instead of libclc, and the driver may reassociate and contract float arithmetic. Results then differ between devices.
narrow_math (bool) – Without libclc installed, evaluate float64 transcendental functions in float32 precision instead of rejecting them.
boundscheck (bool) – Check array indices in this function and raise IndexError from the launch when one is out of bounds. Without it, such accesses read or corrupt unrelated memory of the buffer or are ignored, depending on the driver. The environment variable NUMBA_BOUNDSCHECK=1 turns the check on everywhere.
error_model ({'numpy', 'python'}) – What a division by zero does: 'numpy' (the default, as in numba.cuda) gives NumPy’s result, 'python' raises ZeroDivisionError from the launch, as Numba does on the CPU. The check costs a comparison per division.
narrow (bool, {'ints', 'floats'} or None) – Whether a kernel computes with 32-bit integers and floats where its code says int64 and float64, which Numba uses for Python literals and all index arithmetic. None, the default, narrows integers always, because 64-bit integer arithmetic is slow on GPUs (unless NUMBA_VULKAN_INT64=1), and floats only on devices without float64, with a warning. "ints" narrows only the integers, also on such devices; "floats" narrows only the floats; True narrows both, without a warning; False neither, so the kernel fails on devices without 64-bit types. Arrays of 64-bit elements are converted on the host, and integers that do not fit raise OverflowError. Only the setting of the kernel matters, not that of the functions it calls.
cache (bool) – Keep compiled kernels on disk, so that later processes skip Numba’s compilation, as with Numba’s cache=True. A kernel is compiled again when its source file or that of an @nv.jit function it calls changes. Changes to other functions it uses (@overload, functions of other targets) and to the values of global arrays are not noticed, as in Numba. Functions without a source file (interactive, exec) are not cached.
**options – Accepted for compatibility with Numba’s generic jit and ignored.

Return type:

VulkanDispatcher