numba_vulkan.vectorizers#

@vectorize and @guvectorize for the Vulkan target.

Numba’s decorators look up the class that builds a ufunc in a registry per target; this module registers "vulkan" there, so that

@numba.vectorize(["float32(float32, float32)"], target="vulkan")
def add(a, b):
    return a + b

compiles the scalar function for Vulkan and returns a ufunc-like object that runs it over whole arrays on the device. As with target="cuda", the result is not a real numpy.ufunc: it is called directly, takes NumPy arrays, device arrays and scalars, broadcasts them as NumPy does, and returns a NumPy array, or a device array when any input is one.

Each call runs one kernel, generated from a template, in which every invocation computes one element (for vectorize) or one loop iteration of the core function (for guvectorize). The kernels are compiled per signature and array dimensionality and cached like any other kernel.

Without explicit signatures, the types are taken from the arguments of each call, with Python scalars treated as NumPy 2 does (x * 2.0 keeps the type of a float32 array x).

Attributes#

TARGET

Classes#

`VulkanVectorize`	Builder of an element-wise ufunc, used by `vectorize(target="vulkan")`.
`VulkanUFunc`	An element-wise function over arrays, run on a Vulkan device.
`VulkanGUVectorize`	Builder of a generalized ufunc, used by `guvectorize(target="vulkan")`.
`VulkanGUFunc`	A generalized ufunc run on a Vulkan device.

Module Contents#

numba_vulkan.vectorizers.TARGET = 'vulkan'#

class numba_vulkan.vectorizers.VulkanVectorize(func, identity=None, cache=False, targetoptions=None)#

Builder of an element-wise ufunc, used by vectorize(target="vulkan").

Parameters:

func (function) – The scalar function.
identity (object, optional) – Result of VulkanUFunc.reduce on an empty array.
cache (bool, optional) – Accepted for compatibility; kernels are always cached on disk.
targetoptions (dict, optional) – Options for numba_vulkan.jit, such as fastmath, and device for the default device.

pyfunc#

identity = None#

options#

signatures = []#

add(sig)#

Add a signature such as "float32(float32, float32)".

Parameters:: sig (str or numba.core.typing.Signature) – The signature.

disable_compile()#: Accepted for compatibility: with signatures, only those are used.

build_ufunc()#

Create the ufunc.

Return type:: VulkanUFunc

class numba_vulkan.vectorizers.VulkanUFunc(pyfunc, signatures, options, identity=None)#

An element-wise function over arrays, run on a Vulkan device.

Created by numba.vectorize(..., target="vulkan").

Parameters:

pyfunc (function) – The scalar function.
signatures (list of tuple) – (argument types, return type) pairs; empty to compile for the types of each call.
options (dict) – Options for numba_vulkan.jit, and device.
identity (object, optional) – Result of reduce on an empty array.

Variables:

nout (nin,) – Number of inputs and outputs.
identity (object or None) – As given.

Examples

>>> @numba.vectorize(["float32(float32, float32)"], target="vulkan")
... def hypot(a, b):
...     return math.sqrt(a * a + b * b)
>>> hypot(np.float32([3, 5]), np.float32(4))
array([5.     , 6.40312], dtype=float32)

pyfunc#

identity = None#

nin#

nout = 1#

property types#

The signatures in NumPy’s notation, such as "ff->f".

Return type:: list of str

reduce(array, axis=0, device=None)#

Combine all elements of an array with the function.

Like numpy.ufunc.reduce for a binary function, computed on the device as a tree: pairs of elements are combined in shared memory within each workgroup, and the partial results again until one is left. The order of combination differs from NumPy’s, so the function should be associative (and commutative).

Parameters:

array (array_like or DeviceArray) – The values.
axis ({0, None}) – 0 for a 1-d array; None for all elements of an array of any dimensionality.
device (int, str or Device, optional) – Device to run on. By default that of a device array, or the selected device.

Returns:

The result, a NumPy scalar.

Return type:

scalar

Raises:

TypeError – If the function does not take two arguments, or no signature fits the element type.
ValueError – For an empty array without an identity, or for axis=0 on an array of more than one dimension.

class numba_vulkan.vectorizers.VulkanGUVectorize(func, signature, identity=None, cache=False, targetoptions=None, writable_args=())#

Builder of a generalized ufunc, used by guvectorize(target="vulkan").

Parameters:

func (function) – The core function; it writes its results into its output arguments.
signature (str) – The layout, such as "(n),(n)->()".
identity (optional) – Accepted for compatibility.
cache (optional) – Accepted for compatibility.
writable_args (optional) – Accepted for compatibility.
targetoptions (dict, optional) – Options for numba_vulkan.jit, and device.

pyfunc#

layout#

options#

signatures = []#

add(sig)#

Add a signature such as "void(float32[:], float32[:])".

Parameters:: sig (str or numba.core.typing.Signature) – The signature.

disable_compile()#: Accepted for compatibility: with signatures, only those are used.

build_ufunc()#

Create the generalized ufunc.

Return type:: VulkanGUFunc

class numba_vulkan.vectorizers.VulkanGUFunc(pyfunc, layout, signatures, options)#

A generalized ufunc run on a Vulkan device.

Created by numba.guvectorize(..., target="vulkan"). The core function is called once per element of the loop dimensions, with views of the core dimensions of each argument.

Parameters:

pyfunc (function) – The core function.
layout (str) – The layout, such as "(n),(n)->()".
signatures (list of tuple) – Argument types of each signature; empty to compile for the types of each call, which then requires the outputs to be passed.
options (dict) – Options for numba_vulkan.jit, and device.

Variables:

nout (nin,) – Number of inputs and outputs.

Examples

>>> @numba.guvectorize(["void(float32[:], float32[:], float32[:])"],
...                    "(n),(n)->()", target="vulkan")
... def dot(a, b, out):
...     out[0] = np.dot(a, b)
>>> dot(np.ones((4, 3), np.float32), np.arange(3, dtype=np.float32))
array([3., 3., 3., 3.], dtype=float32)

pyfunc#

layout#