numba_vulkan.compiler#

Compilation pipeline of the Vulkan target.

Attributes#

LOCAL_SIZES

PUSH_CONSTANT_BYTES

Classes#

VulkanCompileResult

Compile result whose entry point is a synthetic identifier.

VulkanBackend

Final compiler pass: package the lowering output in a compile result.

NumberSharedArrays

Give every call that creates an array an identity.

VulkanCompiler

Numba compiler pipeline of the Vulkan target.

Functions#

compile_vulkan(pyfunc, return_type, args[, ...])

Run pyfunc through Numba's pipeline down to LLVM IR.

compile_kernel(cres, ndim[, exact, local_size, aliases])

Wrap a compiled function in a shader entry point and emit SPIR-V.

Module Contents#

numba_vulkan.compiler.LOCAL_SIZES#
class numba_vulkan.compiler.VulkanCompileResult#

Compile result whose entry point is a synthetic identifier.

Numba’s dispatchers and target contexts key compiled functions by their entry point, normally a native address. Nothing is compiled to native code here, so the identity of the result is used instead.

property entry_point#

Unique key of this result.

Return type:

int

class numba_vulkan.compiler.VulkanBackend#

Final compiler pass: package the lowering output in a compile result.

run_pass(state)#

Create the VulkanCompileResult.

Parameters:

state (numba.core.compiler.StateDict) – Compiler state; state.cr is replaced by the result.

Returns:

Always True.

Return type:

bool

class numba_vulkan.compiler.NumberSharedArrays#

Give every call that creates an array an identity.

Shared arrays become workgroup variables, and local arrays (including np.zeros and its relatives) private variables, which must be the same for every execution of the call and different from that of any other call. The pass passes a literal derived from the function and the position of the call as a hidden first argument, which typing turns into the array’s “binding” (see numba_vulkan.vkdecl.SharedArray). Being derived from the source, it is the same in every process, as the kernel cache needs.

run_pass(state)#

Add the hidden argument to the calls.

Parameters:

state (numba.core.compiler.StateDict) – Compiler state.

Returns:

Whether the IR changed.

Return type:

bool

class numba_vulkan.compiler.VulkanCompiler#

Numba compiler pipeline of the Vulkan target.

Numba’s untyped and typed passes are used unchanged, followed by its native lowering and VulkanBackend. No machine code is generated.

define_pipelines()#

Build the pass pipeline.

Return type:

list of numba.core.compiler_machinery.PassManager

numba_vulkan.compiler.compile_vulkan(pyfunc, return_type, args, narrow_math=False, fast_math=False, boundscheck=False, error_model='numpy')#

Run pyfunc through Numba’s pipeline down to LLVM IR.

Parameters:

  • pyfunc (function) – The Python function.

  • return_type (numba.types.Type or None) – Return type, or None to infer it.

  • args (tuple of numba.types.Type) – Argument types.

  • narrow_math (bool) – Compute float64 transcendental functions in float32 when libclc is not available.

  • fast_math (bool) – Use the device’s built-in float32 math functions instead of libclc.

  • boundscheck (bool) – Raise IndexError for array indices that are out of bounds.

  • error_model ({'numpy', 'python'}) – Numba’s error model: 'python' raises ZeroDivisionError for divisions by zero, 'numpy' gives NumPy’s results.

Return type:

VulkanCompileResult

numba_vulkan.compiler.PUSH_CONSTANT_BYTES = 128#
numba_vulkan.compiler.compile_kernel(cres, ndim, exact=True, local_size=None, aliases=())#

Wrap a compiled function in a shader entry point and emit SPIR-V.

Parameters:

  • cres (VulkanCompileResult) – The compiled kernel body; it must return None.

  • ndim (int) – Dimensionality of the dispatch grid (selects the default workgroup size).

  • exact (bool) – Whether drivers must keep float arithmetic as written; see numba_vulkan.codegen.mark_exact.

  • local_size (tuple of int, optional) – Workgroup size; by default 64 invocations, shaped by ndim.

  • aliases (tuple of tuple) – (index, earlier) for array arguments launched with the same buffer as an earlier one: they use its binding, so that the kernel sees that they share memory (as the checks of array expressions assigned to arrays they read need to).

Return type:

CompiledKernel