numba_vulkan.vkimpl#

Lowering of Vulkan-specific functions and of array element access.

Attributes#

`registry`
`lower`

Functions#

`lower_dispatcher_constant`(context, builder, ty, pyval)	Lower a Vulkan dispatcher used as a global or closure variable.
`lower_global_id`(context, builder, sig, args)	Lower `global_id(axis)` to the GlobalInvocationId built-in.
`buffer_element_type`(context, dtype)	LLVM type of one element of a buffer holding `dtype` values.
`half_conversion`(builder, value, target)	Convert between `half` and `float` behind LLVM's back.
`storage_type`(context, aryty)	LLVM type of one element of an array's buffer.
`lower_getitem`(context, builder, sig, args)	Lower `array[index]`.
`lower_getitem_ellipsis`(context, builder, sig, args)	Lower `array[...]`; see lower_getitem.
`lower_transpose`(context, builder, ty, value)	Lower `array.T`: a view with the axes reversed.
`lower_setitem_ellipsis`(context, builder, sig, args)	Lower `array[...] = value`; see lower_setitem.
`lower_setitem`(context, builder, sig, args)	Lower `array[index] = value`.
`lower_getiter`(context, builder, sig, args)	Lower `iter(array)`: an iterator over the first axis.
`lower_iternext`(context, builder, sig, args, result)	Lower `next(iterator)` for an array iterator.
`flat_item`(typingctx, array, position)	Element of an array by its position in row-major order.
`lower_local_size`(context, builder, sig, args)	Lower `local_size(axis)`.
`lower_num_groups`(context, builder, sig, args)	Lower `num_groups(axis)`.
`lower_barrier`(context, builder, sig, args)	Lower `barrier()`; see numba_vulkan.buffers.barrier.
`static_array`(context, builder, aryty, shape)	The value of an array whose shape is known at compile time.
`lower_shared_array`(context, builder, sig, args)	Lower `shared.array(shape, dtype)`.
`lower_local_array`(context, builder, sig, args)	Lower `local.array(shape, dtype)`; see lower_shared_array.
`lower_cas`(context, builder, sig, args)	Lower `atomic.cas(array, index, expected, value)`.
`expression_shape`(context, builder, ty, value)	The broadcast shape of an expression.
`expression_element`(context, builder, ty, value, indices)	Compute one element of an expression.
`lower_expression_getitem`(context, builder, sig, args)	Lower `expression[index]`: compute one element.
`lower_expression_shape`(context, builder, ty, value)	Lower `expression.shape`.
`lower_expression_size`(context, builder, ty, value)	Lower `expression.size`.
`lower_expression_ndim`(context, builder, ty, value)	Lower `expression.ndim`.
`lower_expression_len`(context, builder, sig, args)	Lower `len(expression)`.
`assign_expression`(context, builder, aryty, selection, ...)	Write an expression, element by element, into a view.
`lower_setitem_expression`(context, builder, sig, args)	Lower `array[index] = expression`; see assign_expression.
`lower_print`(context, builder, sig, args)	Lower `print(...)` to a record in the kernel's print buffer.
`lower_broadcast_first`(context, builder, sig, args)	Lower `subgroup.broadcast_first(value)`.
`lower_subgroup_any`(context, builder, sig, args)	Lower `subgroup.any(predicate)`.
`lower_subgroup_all`(context, builder, sig, args)	Lower `subgroup.all(predicate)`.
`lower_subgroup_elect`(context, builder, sig, args)	Lower `subgroup.elect()`.
`lower_subgroup_ballot`(context, builder, sig, args)	Lower `subgroup.ballot(predicate)` to a tuple of four words.
`lower_subgroup_ballot_count`(context, builder, sig, args)	Lower `subgroup.ballot_count(predicate)`.
`lower_record_getattr`(context, builder, ty, value, attr)	Lower `record.field`.
`lower_record_setattr`(context, builder, sig, args, attr)	Lower `record.field = value`.
`lower_record_getitem`(context, builder, sig, args)	Lower `record["field"]`.
`lower_record_setitem`(context, builder, sig, args)	Lower `record["field"] = value`.

Module Contents#

numba_vulkan.vkimpl.registry#

numba_vulkan.vkimpl.lower#

numba_vulkan.vkimpl.lower_dispatcher_constant(context, builder, ty, pyval)#

Lower a Vulkan dispatcher used as a global or closure variable.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
ty (VulkanDispatcherType) – Type of the constant.
pyval (VulkanDispatcher) – The dispatcher.

Returns:

A dummy value; calls are resolved at compile time.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.lower_global_id(context, builder, sig, args)#

Lower global_id(axis) to the GlobalInvocationId built-in.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
sig (numba.core.typing.Signature) – Signature the call was typed with.
args (sequence of llvmlite.ir.Value) – Argument values, in the types of sig.args.

Returns:

The int32 invocation index along the literal axis.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.buffer_element_type(context, dtype)#

LLVM type of one element of a buffer holding dtype values.

Parameters:

context (VulkanTargetContext) – The target context.
dtype (numba.types.Type) – Numba element type.

Returns:

i32 for booleans (SPIR-V has no storable bool), otherwise the data type of the integer or float.

Return type:

llvmlite.ir.Type

Raises:

VulkanUnsupportedError – If dtype is not a boolean, integer or float.

numba_vulkan.vkimpl.half_conversion(builder, value, target)#

Convert between half and float behind LLVM’s back.

LLVM would fold fptrunc(fpext(x) * 2) into a half multiply, which needs the shaderFloat16 device feature. The conversions are therefore placeholder calls until after optimisation (see numba_vulkan.buffers.expand_buffer_access).

Parameters:

builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
value (llvmlite.ir.Value) – A half or float value.
target (llvmlite.ir.Type) – The other of the two types.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.storage_type(context, aryty)#

LLVM type of one element of an array’s buffer.

Parameters:

context (VulkanTargetContext) – The target context.
aryty (VulkanArray) – The array type.

Returns:

half for arrays of float16 values, otherwise see buffer_element_type.

Return type:

llvmlite.ir.Type

numba_vulkan.vkimpl.lower_getitem(context, builder, sig, args)#

Lower array[index].

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
sig (numba.core.typing.Signature) – Signature the call was typed with.
args (sequence of llvmlite.ir.Value) – Argument values, in the types of sig.args.

Returns:

The element, or a view of the same buffer if the index contains slices or covers only the leading axes.

Return type:

llvmlite.ir.Value

Raises:

VulkanUnsupportedError – If the index contains anything but integers and slices.

numba_vulkan.vkimpl.lower_getitem_ellipsis(context, builder, sig, args)#: Lower array[...]; see lower_getitem.

numba_vulkan.vkimpl.lower_transpose(context, builder, ty, value)#

Lower array.T: a view with the axes reversed.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
ty (VulkanArray) – Type of the array.
value (llvmlite.ir.Value) – The array.

Returns:

The transposed view; the array itself for fewer than two axes.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.lower_setitem_ellipsis(context, builder, sig, args)#: Lower array[...] = value; see lower_setitem.

numba_vulkan.vkimpl.lower_setitem(context, builder, sig, args)#

Lower array[index] = value.

If the index selects one element, the value is cast to the element type and stored. If it selects a view, the view is filled with a scalar, or with the elements of an array of the same shape.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
sig (numba.core.typing.Signature) – Signature the call was typed with.
args (sequence of llvmlite.ir.Value) – Argument values, in the types of sig.args.

Returns:

A dummy value.

Return type:

llvmlite.ir.Value

Raises:

VulkanUnsupportedError – If the index contains anything but integers and slices, if the value is an array of another dimensionality, or if it is an array in the same buffer as the target (the copy could overlap).

numba_vulkan.vkimpl.lower_getiter(context, builder, sig, args)#

Lower iter(array): an iterator over the first axis.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
sig (numba.core.typing.Signature) – Signature the call was typed with.
args (sequence of llvmlite.ir.Value) – The array.

Returns:

The iterator.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.lower_iternext(context, builder, sig, args, result)#

Lower next(iterator) for an array iterator.

Yields elements of a 1-d array and views along the first axis otherwise.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
sig (numba.core.typing.Signature) – Signature the call was typed with.
args (sequence of llvmlite.ir.Value) – The iterator.
result (numba.core.imputils._IternextResult) – Receives the validity flag and the yielded value.

numba_vulkan.vkimpl.flat_item(typingctx, array, position)#

Element of an array by its position in row-major order.

Parameters:

array (VulkanArray or VulkanExpr) – The array or array expression, of any dimensionality.
position (int) – Position between 0 and array.size - 1; it is not checked.

Returns:

The element.

Return type:

scalar

numba_vulkan.vkimpl.lower_local_size(context, builder, sig, args)#

Lower local_size(axis).

The workgroup size is only fixed when the kernel is compiled, after its functions; a placeholder call stands for it until then (see numba_vulkan.codegen.VulkanCodeLibrary).

Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_num_groups(context, builder, sig, args)#

Lower num_groups(axis).

Grids larger than the device allows are dispatched in parts, for which the NumWorkgroups built-in gives the size of the part. The size of the whole grid is passed as push constants instead; a placeholder call stands for it until the push-constant block is laid out (see numba_vulkan.codegen.num_groups_members).

Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_barrier(context, builder, sig, args)#

Lower barrier(); see numba_vulkan.buffers.barrier.

Returns:: A dummy value.
Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.static_array(context, builder, aryty, shape)#

The value of an array whose shape is known at compile time.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
aryty (VulkanArray) – Type of the array.
shape (tuple of int) – Its shape.

Returns:

The metadata structure, with C-contiguous strides and offset 0.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.lower_shared_array(context, builder, sig, args)#

Lower shared.array(shape, dtype).

The memory is a workgroup variable named after the array’s binding, declared when buffer accesses are expanded; here only the metadata of the array is built.

Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_local_array(context, builder, sig, args)#

Lower local.array(shape, dtype); see lower_shared_array.

Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_cas(context, builder, sig, args)#

Lower atomic.cas(array, index, expected, value).

Returns:: The element’s previous value.
Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.expression_shape(context, builder, ty, value)#

The broadcast shape of an expression.

Raises ValueError in the kernel if the operands cannot be broadcast together.

Returns:: One intp extent per dimension.
Return type:: list of llvmlite.ir.Value

numba_vulkan.vkimpl.expression_element(context, builder, ty, value, indices)#

Compute one element of an expression.

Parameters:

context (VulkanTargetContext) – The target context.
builder (llvmlite.ir.IRBuilder) – Builder positioned where the code is emitted.
ty (VulkanExpr) – Type of the expression.
value (llvmlite.ir.Value) – The expression.
indices (list of llvmlite.ir.Value) – One non-negative, in-range intp index per dimension of ty.

Returns:

The element, of type ty.dtype.

Return type:

llvmlite.ir.Value

numba_vulkan.vkimpl.lower_expression_getitem(context, builder, sig, args)#

Lower expression[index]: compute one element.

Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_expression_shape(context, builder, ty, value)#: Lower expression.shape.

numba_vulkan.vkimpl.lower_expression_size(context, builder, ty, value)#: Lower expression.size.

numba_vulkan.vkimpl.lower_expression_ndim(context, builder, ty, value)#: Lower expression.ndim.

numba_vulkan.vkimpl.lower_expression_len(context, builder, sig, args)#: Lower len(expression).

numba_vulkan.vkimpl.assign_expression(context, builder, aryty, selection, exprty, expr)#

Write an expression, element by element, into a view.

Parameters:

aryty (VulkanArray) – Type of the array written to.
selection (_Selection) – The view of it that is written.
exprty (VulkanExpr) – Type of the expression.
expr (llvmlite.ir.Value) – The expression.

Notes

NumPy computes the right-hand side before writing. Element by element, that is the same only if the expression reads the written array at the very positions it writes, or only outside the span of positions it writes; anything else raises ValueError.

numba_vulkan.vkimpl.lower_setitem_expression(context, builder, sig, args)#

Lower array[index] = expression; see assign_expression.

Returns:: A dummy value.
Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_print(context, builder, sig, args)#

Lower print(...) to a record in the kernel’s print buffer.

The buffer starts with a cursor and its capacity, in 32-bit words. A call reserves room for its record with an atomic addition to the cursor and, if it fits, writes the number of its format followed by the values. The host prints the records after the kernel (see numba_vulkan.runtime.print_records). String constants are part of the format and are not written.

Returns:: A dummy value.
Return type:: llvmlite.ir.Value

numba_vulkan.vkimpl.lower_broadcast_first(context, builder, sig, args)#: Lower subgroup.broadcast_first(value).

numba_vulkan.vkimpl.lower_subgroup_any(context, builder, sig, args)#: Lower subgroup.any(predicate).

numba_vulkan.vkimpl.lower_subgroup_all(context, builder, sig, args)#: Lower subgroup.all(predicate).

numba_vulkan.vkimpl.lower_subgroup_elect(context, builder, sig, args)#: Lower subgroup.elect().

numba_vulkan.vkimpl.lower_subgroup_ballot(context, builder, sig, args)#: Lower subgroup.ballot(predicate) to a tuple of four words.

numba_vulkan.vkimpl.lower_subgroup_ballot_count(context, builder, sig, args)#: Lower subgroup.ballot_count(predicate).

numba_vulkan.vkimpl.lower_record_getattr(context, builder, ty, value, attr)#: Lower record.field.

numba_vulkan.vkimpl.lower_record_setattr(context, builder, sig, args, attr)#: Lower record.field = value.

numba_vulkan.vkimpl.lower_record_getitem(context, builder, sig, args)#: Lower record["field"].

numba_vulkan.vkimpl.lower_record_setitem(context, builder, sig, args)#: Lower record["field"] = value.