numba_vulkan.buffers#

Access to Vulkan storage buffers from LLVM IR.

Shaders have no general pointers, and LLVM’s optimiser readily creates pointer phis and selects that SPIR-V cannot express. Buffer accesses are therefore emitted as calls to opaque placeholder functions that take a binding and an element index, and only expand_buffer_access() turns them into real loads and stores, after all optimisation has happened.

Attributes#

i32

META_BINDING

STATUS_INDEX

PRINT_BINDING

PUSH_BINDING

PUSH_BLOCK

print_formats

CONSTANT_BASE

SHARED_BASE

LOCAL_BASE

shared_sizes

ATOMIC_OPS

Functions#

renumber_print(text, binding)

Give the print buffer of a kernel its binding.

constant_binding(array)

Binding that stands for a NumPy array used as a global constant.

constant_order(text)

The constant arrays that LLVM IR refers to, in order of appearance.

constant_data(binding)

Contents of the constant array behind a placeholder binding.

renumber_constants(text, first[, order])

Give the constant arrays of a kernel bindings after its arguments.

arg_binding(index)

Descriptor binding of a kernel argument.

load_element(builder, binding, elem, index)

Read one element of a storage buffer.

store_element(builder, binding, elem, index, value)

Write one element of a storage buffer.

atomic_element(builder, binding, elem, index, op, value)

Apply an atomic read-modify-write to one element of a buffer.

compare_and_swap(builder, binding, index, expected, value)

Atomically replace one 32-bit element if it has the expected value.

barrier(builder)

Wait until all invocations of the workgroup arrive here.

expand_buffer_access(text[, push_types])

Replace the placeholder calls in LLVM IR by real buffer accesses.

Module Contents#

numba_vulkan.buffers.i32#
numba_vulkan.buffers.META_BINDING = 0#
numba_vulkan.buffers.STATUS_INDEX = 0#
numba_vulkan.buffers.PRINT_BINDING = 1048575#
numba_vulkan.buffers.PUSH_BINDING = 1048574#
numba_vulkan.buffers.PUSH_BLOCK = 'nv.args'#
numba_vulkan.buffers.print_formats#
numba_vulkan.buffers.renumber_print(text, binding)#

Give the print buffer of a kernel its binding.

Parameters:

  • text (str) – Textual LLVM IR with placeholder accesses.

  • binding (int) – The binding to use.

Returns:

  • text (str) – The IR with the placeholder binding replaced.

  • used (bool) – Whether the kernel prints.

numba_vulkan.buffers.CONSTANT_BASE = 1048576#
numba_vulkan.buffers.SHARED_BASE = 16777216#
numba_vulkan.buffers.LOCAL_BASE = 33554432#
numba_vulkan.buffers.shared_sizes#
numba_vulkan.buffers.constant_binding(array)#

Binding that stands for a NumPy array used as a global constant.

The array is copied when it is first seen: like Numba on the CPU, the compiled code keeps the values the array had at that time.

Parameters:

array (numpy.ndarray) – The array.

Returns:

A binding at or above CONSTANT_BASE, the same for every use of the same array object.

Return type:

int

numba_vulkan.buffers.constant_order(text)#

The constant arrays that LLVM IR refers to, in order of appearance.

Parameters:

text (str) – Textual LLVM IR with placeholder buffer accesses.

Returns:

The placeholder bindings (at or above CONSTANT_BASE), each once.

Return type:

list of int

numba_vulkan.buffers.constant_data(binding)#

Contents of the constant array behind a placeholder binding.

Parameters:

binding (int) – A binding returned by constant_binding.

Returns:

The values the array had when it was first seen.

Return type:

numpy.ndarray

numba_vulkan.buffers.renumber_constants(text, first, order=None)#

Give the constant arrays of a kernel bindings after its arguments.

Parameters:

  • text (str) – Textual LLVM IR with placeholder buffer accesses.

  • first (int) – Binding for the first constant array.

  • order (list of int, optional) – Placeholder bindings in the order in which they receive their bindings; those that text no longer uses are skipped. By default the order of appearance in text.

Returns:

  • text (str) – The IR with the placeholder bindings replaced.

  • placeholders (dict of int to int) – Placeholder binding of each constant array by its binding in this kernel.

numba_vulkan.buffers.arg_binding(index)#

Descriptor binding of a kernel argument.

Parameters:

index (int) – Position of the argument in the kernel signature.

Returns:

1 + index; binding 0 holds the error status and the array shapes.

Return type:

int

numba_vulkan.buffers.load_element(builder, binding, elem, index)#

Read one element of a storage buffer.

Emits a call to a placeholder function, which expand_buffer_access later turns into the actual load.

Parameters:

  • builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.

  • binding (int or llvmlite.ir.Value) – Descriptor binding of the buffer: a number, or a value that is a constant once the kernel is inlined.

  • elem (llvmlite.ir.Type) – Element type of the buffer.

  • index (llvmlite.ir.Value) – Element index.

Returns:

The element, of type elem.

Return type:

llvmlite.ir.Value

numba_vulkan.buffers.store_element(builder, binding, elem, index, value)#

Write one element of a storage buffer.

Emits a call to a placeholder function, which expand_buffer_access later turns into the actual store.

Parameters:

  • builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.

  • binding (int or llvmlite.ir.Value) – Descriptor binding of the buffer: a number, or a value that is a constant once the kernel is inlined.

  • elem (llvmlite.ir.Type) – Element type of the buffer.

  • index (llvmlite.ir.Value) – Element index.

  • value (llvmlite.ir.Value) – Value to store, of type elem.

numba_vulkan.buffers.ATOMIC_OPS = ('add', 'sub', 'and', 'or', 'xor', 'xchg', 'max', 'min', 'umax', 'umin', 'fadd')#
numba_vulkan.buffers.atomic_element(builder, binding, elem, index, op, value)#

Apply an atomic read-modify-write to one element of a buffer.

Emits a call to a placeholder function, which expand_buffer_access later turns into an atomicrmw instruction with relaxed ordering.

Parameters:

  • builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.

  • binding (int or llvmlite.ir.Value) – Descriptor binding of the buffer (see load_element), or the binding of a shared array.

  • elem (llvmlite.ir.Type) – Element type: a 32-bit integer, or float for fadd.

  • index (llvmlite.ir.Value) – Element index.

  • op (str) – One of ATOMIC_OPS.

  • value (llvmlite.ir.Value) – The operand, of type elem.

Returns:

The value the element had before.

Return type:

llvmlite.ir.Value

numba_vulkan.buffers.compare_and_swap(builder, binding, index, expected, value)#

Atomically replace one 32-bit element if it has the expected value.

Emits a call to a placeholder function, which expand_buffer_access later expands.

Parameters:

  • builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.

  • binding (int or llvmlite.ir.Value) – Descriptor binding of the buffer (see load_element), or the binding of a shared array.

  • index (llvmlite.ir.Value) – Element index.

  • expected (llvmlite.ir.Value) – i32 values: what the element must hold, and its replacement.

  • value (llvmlite.ir.Value) – i32 values: what the element must hold, and its replacement.

Returns:

The value the element had before; the swap happened if it equals expected.

Return type:

llvmlite.ir.Value

numba_vulkan.buffers.barrier(builder)#

Wait until all invocations of the workgroup arrive here.

Emits a call to a placeholder function. It has unknown memory effects, so LLVM moves no memory access across it, and it is convergent, so LLVM does not duplicate it into different branches.

Parameters:

builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.

numba_vulkan.buffers.expand_buffer_access(text, push_types=())#

Replace the placeholder calls in LLVM IR by real buffer accesses.

Each placeholder becomes a llvm.spv.resource.handlefrombinding call, a llvm.spv.resource.getpointer call and a load or store. Loads from PUSH_BINDING become loads of members of the push-constant block. The declarations and name strings these need are appended.

Parameters:

  • text (str) – Textual LLVM IR after optimisation.

  • push_types (sequence of str, optional) – LLVM types of the members of the push-constant block, if the kernel has one.

Returns:

  • text (str) – The rewritten LLVM IR.

  • written (set of int) – Bindings the code stores to.

  • read (set of int) – Bindings the code loads from.

Raises:

SpirvCodegenError – If a placeholder call with a non-constant binding remains.

Notes

This works on text because llvmlite cannot build IR that uses LLVM target extension types, and because it has to run after LLVM’s optimisation passes.