numba_vulkan.runtime#

Minimal Vulkan compute runtime: devices, device arrays and kernel launch.

Attributes#

Classes#

DeviceInfo

A Vulkan physical device.

Device

A logical device with one compute queue and a pipeline cache.

DeviceArray

An array that lives on a Vulkan device.

Stream

An ordered sequence of copies and launches on a device.

Event

A point in the work submitted to a device, like a CUDA event.

Functions#

debug_level()

How much NUMBA_VULKAN_DEBUG asks Vulkan's validation layer to check.

debug_enabled()

Whether NUMBA_VULKAN_DEBUG asks for Vulkan's validation layer.

validation_messages([clear])

What the validation layer reported so far (NUMBA_VULKAN_DEBUG=1).

list_devices([refresh])

All Vulkan 1.2+ devices that expose a compute queue.

print_records(words)

Print what a kernel's print calls recorded.

to_device(array[, device, stream])

Copy a NumPy array to a device.

device_array(shape[, dtype, device])

Create an array on a device without initialising it.

device_array_like(array[, device])

Create an uninitialised device array with the shape and type of another.

get_device([which])

Return a Device.

stream([device])

Create a Stream on a device, like numba.cuda.stream.

pinned_array(shape[, dtype, device])

Allocate a host array that a device can copy to and from directly.

event([device])

Create an Event on a device.

event_elapsed_time(start, end)

Milliseconds between two events, as numba.cuda names it.

synchronize([device])

Wait until launched kernels have finished.

select_device(which)

Make which the default device for subsequent kernel calls.

Module Contents#

numba_vulkan.runtime.hint = " Install your distribution's Vulkan loader (libvulkan1 on Debian and Ubuntu) and a driver."#
class numba_vulkan.runtime.DeviceInfo#

A Vulkan physical device.

Variables:
  • index (int) – Position in list_devices.

  • name (str) – Device name reported by the driver.

  • kind ({'discrete', 'integrated', 'virtual', 'cpu', 'other'}) – Device type.

  • int8 (float64, int64, int16,) – Whether shaders may use the respective type.

  • handle (object) – The VkPhysicalDevice handle.

  • int64_atomics (bool) – Whether shaders may apply atomic operations to 64-bit integers.

  • float16 (bool) – Whether shaders may compute with float16 values.

  • storage16 (bool) – Whether buffers may hold 16-bit values (float16 arrays).

  • storage8 (bool) – Whether buffers may hold 8-bit values (int8 and uint8 arrays).

  • float32_atomic_add (bool) – Whether shaders may add to float32 values atomically, in buffers and in shared memory (VK_EXT_shader_atomic_float).

  • max_local_size (tuple of int) – Largest workgroup extent along each axis.

  • max_local_invocations (int) – Largest number of invocations in one workgroup.

  • max_groups (tuple of int) – Largest number of workgroups along each axis of a dispatch.

  • max_shared_memory (int) – Bytes of workgroup-shared memory available to a kernel.

  • nan_preserve (tuple of int) – Float widths for which shaders can require NaN, infinity and signed zero to be kept (SignedZeroInfNanPreserve).

  • timestamp_bits (int) – Valid bits of the timestamps the compute queue writes; 0 if it writes none (see event).

  • timestamp_period (float) – Nanoseconds per timestamp tick.

  • subgroup_size (int) – Invocations per subgroup (subgroup.size() in kernels).

  • subgroup_shuffle_relative (subgroup_basic, subgroup_vote, subgroup_arithmetic, subgroup_ballot, subgroup_shuffle,) – Which classes of subgroup operations compute shaders may use; see numba_vulkan.stubs.subgroup.

index: int#
name: str#
kind: str#
float64: bool#
int64: bool#
int16: bool#
int8: bool#
handle: object#
int64_atomics: bool = False#
float16: bool = False#
storage16: bool = False#
storage8: bool = False#
float32_atomic_add: bool = False#
max_local_size: tuple = (128, 128, 64)#
max_local_invocations: int = 128#
max_groups: tuple = (65535, 65535, 65535)#
max_shared_memory: int = 16384#
nan_preserve: tuple = ()#
timestamp_bits: int = 0#
timestamp_period: float = 1.0#
subgroup_size: int = 1#
subgroup_basic: bool = False#
subgroup_vote: bool = False#
subgroup_arithmetic: bool = False#
subgroup_ballot: bool = False#
subgroup_shuffle: bool = False#
subgroup_shuffle_relative: bool = False#
numba_vulkan.runtime.DEBUG_ENV_VAR = 'NUMBA_VULKAN_DEBUG'#
numba_vulkan.runtime.debug_level()#

How much NUMBA_VULKAN_DEBUG asks Vulkan’s validation layer to check.

Returns:

0: nothing (the default); 1: the core checks of API usage; 2: also synchronization, that is missing barriers between commands, which is slower.

Return type:

int

numba_vulkan.runtime.debug_enabled()#

Whether NUMBA_VULKAN_DEBUG asks for Vulkan’s validation layer.

Return type:

bool

numba_vulkan.runtime.validation_messages(clear=False)#

What the validation layer reported so far (NUMBA_VULKAN_DEBUG=1).

Parameters:

clear (bool, optional) – Forget the messages after returning them.

Returns:

(severity, message) pairs, severity "warning" or "error", oldest first.

Return type:

list of tuple

numba_vulkan.runtime.list_devices(refresh=False)#

All Vulkan 1.2+ devices that expose a compute queue.

Parameters:

refresh (bool) – Ask the Vulkan loader again instead of returning the devices found by the first call.

Return type:

list of DeviceInfo

class numba_vulkan.runtime.Device(info)#

A logical device with one compute queue and a pipeline cache.

Parameters:

info (DeviceInfo) – The physical device to open.

Variables:
  • info (DeviceInfo) – The physical device.

  • handle (object) – The VkDevice handle.

  • queue (object) – The compute queue.

  • mode (numba_vulkan.narrowing.Mode) – How kernels are compiled for the device by default: with 32-bit integers, and with 32-bit floats if it lacks float64.

  • pool_limit (int) – Number of bytes of released buffers kept for reuse; set initially from the environment variable NUMBA_VULKAN_POOL_MB (default 1024).

Notes

A device is not thread-safe: launch kernels and transfer arrays from one thread at a time.

info#
family#
transfer_family#
handle#
queue#
stream_queue#
memory#
command_pool#
mode#
pool_limit#
check_support(kernel)#

Check that this device can run a kernel.

Parameters:

kernel (numba_vulkan.codegen.CompiledKernel) – The compiled kernel.

Raises:

VulkanSupportError – If the device lacks a capability the kernel needs.

trim(limit=0)#

Free pooled buffers until at most limit bytes stay pooled.

Buffers of arrays that are no longer in use are kept for reuse, up to pool_limit bytes; this releases them to the system.

Parameters:

limit (int) – Number of bytes that may remain pooled.

MAX_BATCH = 16#
MAX_IN_FLIGHT = 4#
MAX_SLOTS = 32#
launch(kernel, groups, arrays, push=b'')#

Start kernel over groups workgroups without waiting for it.

Used for launches whose arrays are all device arrays, so that nothing has to be read back. Launches are recorded into a command buffer, which is submitted at once if the device has nothing else to do, and otherwise when it holds MAX_BATCH launches or at the next synchronisation, whichever comes first. A barrier precedes each launch, so it sees the writes of earlier ones.

A launch uses a slot of the kernel: a descriptor set and buffers for its small host arrays. Launches with the same buffers and the same host arrays share one, so repeated launches neither write descriptors nor copy data; their scalars and shapes are push constants, recorded with each launch.

An exception raised by the kernel is reported by the next synchronize, which happens before any copy from or to the device and before any launch that waits.

Parameters:
  • kernel (CompiledKernel) – The compiled kernel.

  • groups (tuple of int) – Workgroup counts along x, y and z.

  • arrays (list of numpy.ndarray, DeviceArray or None) – One per binding: device arrays, and small host arrays (the status and, for kernels without push constants, shapes and scalars). None for the bindings in kernel.unbound.

  • push (bytes, optional) – The kernel’s push constants, kernel.push_size bytes.

Raises:

ValueError – If a device array belongs to another device.

synchronize(report=True)#

Wait until every kernel launched on this device has finished.

Parameters:

report (bool) – Whether to raise the exception of a kernel that raised one since the last synchronisation.

Raises:

Exception – The first exception that a kernel launched asynchronously raised since then; see numba_vulkan.dispatcher.

wait_for(buffer)#

Wait until no launch uses a buffer any more.

Launches that do not use it may go on running; with memory that the host can map, copies to and from the buffer then overlap with them.

Parameters:

buffer (_Buffer) – The buffer.

Raises:

Exception – As synchronize does.

launch_stream(stream, kernel, groups, arrays, push=b'')#

Enqueue kernel on a Stream; see launch for the arguments.

The arrays must be device arrays (and the host arrays that hold the status and, without push constants, shapes and scalars). What the kernel prints is printed once the stream has passed the launch.

run(kernel, groups, arrays, push=b'')#

Dispatch kernel over groups workgroups.

Parameters:
  • kernel (CompiledKernel) – The compiled kernel.

  • groups (tuple of int) – Workgroup counts along x, y and z.

  • arrays (list of numpy.ndarray, DeviceArray or None) – One array per binding. C-contiguous host arrays are uploaded, and those the kernel writes to are copied back afterwards; device arrays are used in place. None for the bindings in kernel.unbound.

  • push (bytes, optional) – The kernel’s push constants, kernel.push_size bytes.

Raises:

ValueError – If a device array belongs to another device.

numba_vulkan.runtime.PRINT_BUFFER_WORDS#
numba_vulkan.runtime.print_records(words)#

Print what a kernel’s print calls recorded.

Parameters:

words (numpy.ndarray) – The print buffer as uint32: cursor, capacity, then records; see numba_vulkan.vkimpl.lower_print.

class numba_vulkan.runtime.DeviceArray(device, shape, dtype, _stored=None)#

An array that lives on a Vulkan device.

Kernels use device arrays in place, without the copies to and from the device that NumPy arrays need on every call. Create them with to_device, device_array or device_array_like.

Indexing works as on NumPy arrays, without advanced indexing: an integer for every axis reads or writes one element, and slices, integers for some axes, ... and None give a view that shares the memory, which kernels accept like any device array. .T, transpose and reshape give views as well (reshape only where the strides allow, as NumPy without copying).

Parameters:
  • device (Device) – The device that holds the array.

  • shape (tuple of int) – Shape of the array.

  • dtype (numpy.dtype) – Element type. Boolean arrays are stored as int32 on the device, and 64-bit types as their 32-bit counterparts on devices that lack them (see numba_vulkan.narrowing).

Variables:
  • device (Device) – The device that holds the array.

  • shape (tuple of int) – Shape of the array.

  • dtype (numpy.dtype) – Element type.

Examples

>>> x = nv.to_device(np.arange(8, dtype=np.float32))
>>> out = nv.device_array_like(x)
>>> kernel.forall(8)(x, out)
>>> out.copy_to_host()
>>> out[2:6].copy_to_host()     # a view of four elements
>>> out[0] = 1                  # writes one element
device#
shape#
dtype#
property ndim#

Number of dimensions.

property size#

Number of elements.

property nbytes#

Size of the elements in bytes, as they are stored on the host.

property strides#

Steps between elements along each axis in bytes, as in NumPy.

property is_contiguous#

Whether the elements lie one after the other in C order.

Return type:

bool

property T#

The array with its axes reversed, as a view.

transpose(*axes)#

Permute the axes, as a view.

Parameters:

*axes (int or tuple of int) – The new order of the axes; reversed by default.

Return type:

DeviceArray

reshape(*shape)#

Give the array a new shape without copying.

Parameters:

*shape (int or tuple of int) – The new shape, with as many elements as the array has; one extent may be -1, as in NumPy.

Returns:

An array that shares the data with this one.

Return type:

DeviceArray

Raises:

ValueError – If the number of elements differs, or if the layout of the array does not allow the shape without copying; copy first.

ravel()#

The array as one dimension, a view where the layout allows it.

Return type:

DeviceArray

copy()#

A contiguous copy on the same device.

Return type:

DeviceArray

copy_to_device(array, stream=None)#

Overwrite the contents with those of a host array.

Parameters:
  • array (array_like) – Values of the same shape; converted to the element type.

  • stream (Stream, optional) – Enqueue the copy on a stream and return at once; see Stream. The array (if it is not pinned, its converted copy) must not change before the stream has done the copy.

Returns:

The array itself.

Return type:

DeviceArray

Raises:
  • ValueError – If the shapes differ.

  • Exception – One that a kernel launched asynchronously raised; see Device.synchronize.

copy_to_host(out=None, stream=None)#

Copy the contents to a NumPy array.

Parameters:
  • out (numpy.ndarray, optional) – C-contiguous array of the same shape and element type to fill. A new array is created if omitted.

  • stream (Stream, optional) – Enqueue the copy on a stream and return at once; out holds the data after stream.synchronize().

Return type:

numpy.ndarray

Raises:
  • ValueError – If out does not match the array.

  • Exception – One that a kernel launched asynchronously raised; see Device.synchronize.

sum()#

Sum of all elements, computed on the device.

Return type:

scalar

min()#

Smallest element, computed on the device.

Return type:

scalar

max()#

Largest element, computed on the device.

Return type:

scalar

numba_vulkan.runtime.to_device(array, device=None, stream=None)#

Copy a NumPy array to a device.

Parameters:
  • array (array_like) – The values.

  • device (int, str, DeviceInfo, Device or None) – The device; see get_device. Defaults to the selected device, or that of stream.

  • stream (Stream, optional) – Enqueue the copy on a stream and return at once; see Stream.

Return type:

DeviceArray

numba_vulkan.runtime.device_array(shape, dtype=np.float64, device=None)#

Create an array on a device without initialising it.

Parameters:
Return type:

DeviceArray

numba_vulkan.runtime.device_array_like(array, device=None)#

Create an uninitialised device array with the shape and type of another.

Parameters:
Return type:

DeviceArray

numba_vulkan.runtime.get_device(which=None)#

Return a Device.

Parameters:

which (int, str, DeviceInfo, Device or None) – Index into list_devices(), a case-insensitive substring of the device name, or None for the current default.

Return type:

Device

class numba_vulkan.runtime.Stream(device)#

An ordered sequence of copies and launches on a device.

Work on different streams can run at the same time: copies to and from the device on a copy engine, kernels on the device, as with streams in numba.cuda. Work on one stream runs in the order it was enqueued. Create streams with stream, and pass them as stream= to to_device, DeviceArray.copy_to_device and DeviceArray.copy_to_host, and as the third element of kernel[groups, local_size, stream] or as forall(n, stream=...).

All of these return at once. Data copied to the host is there after synchronize; an exception a kernel raised is raised by it. Copies from and to arrays created with pinned_array go directly between that memory and the device; other host arrays go through a staging buffer, and the data of a copy to the host is copied out of it by synchronize. NumPy arrays passed to a launch on a stream are copied to the device before it and, if the kernel writes them, back after it; what a kernel prints is printed by synchronize as well. Views with gaps are copied run by run of adjacent elements, or with many runs by a kernel that gathers or scatters them on the device.

Each stream has a timeline semaphore: every submission waits for the previous one of its stream, and for the streams that last used its buffers. Copies go to a queue of a transfer-only queue family where the device has one, kernels to a second compute queue where there is one. Work enqueued without a stream waits on the host for streams that still use its buffers, and the other way round.

Parameters:

device (Device) – The device.

device#
value = 0#
query()#

Whether all work enqueued on the stream has finished.

Return type:

bool

synchronize()#

Wait until all work enqueued on the stream has finished.

Raises:

Exception – The first exception that a kernel launched on the stream raised since the last synchronisation.

auto_synchronize()#

Wait for the stream when the block ends, as in numba.cuda.

Yields:

Stream – The stream itself.

numba_vulkan.runtime.stream(device=None)#

Create a Stream on a device, like numba.cuda.stream.

Parameters:

device (int, str, Device or None) – The device; the selected one by default.

Return type:

Stream

Examples

>>> s = nv.stream()
>>> d = nv.to_device(chunk, stream=s)        # returns at once
>>> kernel[groups, 64, s](d, out)            # after the copy, on the device
>>> host = out.copy_to_host(stream=s)        # valid after synchronize
>>> s.synchronize()
numba_vulkan.runtime.pinned_array(shape, dtype=np.float64, device=None)#

Allocate a host array that a device can copy to and from directly.

Like numba.cuda.pinned_array: the memory belongs to the device and is mapped into the process, so copies on a Stream need no staging buffer and run while the host goes on. Use it for the host side of data streamed in chunks.

Parameters:
  • shape (int or tuple of int) – Shape of the array.

  • dtype (numpy.dtype) – Element type.

  • device (int, str, Device or None) – The device; the selected one by default.

Returns:

Uninitialised, C-contiguous.

Return type:

numpy.ndarray

class numba_vulkan.runtime.Event(device)#

A point in the work submitted to a device, like a CUDA event.

record marks the point after all work launched so far; the device writes the time at which it gets there. elapsed_time gives the time between two events on the device, which excludes the cost of launching and waiting in Python, as long as the device has work: as with CUDA’s events, time in which it waits for the host between two events counts. On a Stream, record(stream) marks the point after the work enqueued on that stream instead. Create events with event.

Parameters:

device (Device) – The device.

Raises:

VulkanSupportError – If the device’s compute queue writes no timestamps.

QUERIES = 1024#
device#
record(stream=None)#

Mark the point after all work launched on the device so far.

Parameters:

stream (Stream, optional) – Mark the point after the work enqueued on this stream so far instead, like cuda.event().record(stream).

Returns:

The event itself.

Return type:

Event

Raises:

ValueError – If the stream belongs to another device.

synchronize()#

Wait until the device has passed the event.

elapsed_time(end)#

Milliseconds the device took from this event to end.

Parameters:

end (Event) – A later event on the same device.

Return type:

float

numba_vulkan.runtime.event(device=None)#

Create an Event on a device.

Parameters:

device (int, str, Device or None) – The device; the selected one by default.

Return type:

Event

Examples

>>> start, end = nv.event(), nv.event()
>>> start.record()
>>> kernel.forall(n)(x, out)      # device arrays: does not wait
>>> end.record()
>>> start.elapsed_time(end)       # milliseconds on the device
numba_vulkan.runtime.event_elapsed_time(start, end)#

Milliseconds between two events, as numba.cuda names it.

Return type:

float

numba_vulkan.runtime.synchronize(device=None)#

Wait until launched kernels have finished.

Launches whose arrays are all device arrays return before the kernel has run, as in numba.cuda. Copying a device array waits on its own, so this is mainly needed to measure time.

Parameters:

device (int, str, DeviceInfo, Device or None) – The device to wait for; every device that was used if None.

numba_vulkan.runtime.select_device(which)#

Make which the default device for subsequent kernel calls.

Parameters:

which (int, str or DeviceInfo) – See get_device().

Return type:

Device