numba_vulkan.codegen#

Code libraries for the Vulkan target: LLVM IR in, SPIR-V out.

Attributes#

`TRIPLE`
`ENTRY_POINT`
`PASSES`
`EARLY_PASSES`
`EMIT_TIMEOUT`
`LLC_ENV_VAR`
`emitter`
`LOCAL_SIZE_ARRAY`
`GROUP_PREFIX`

Classes#

`Emitter`	A child process that runs LLVM's SPIR-V backend.
`CompiledKernel`	A kernel specialisation ready to be run by numba_vulkan.runtime.
`VulkanCodeLibrary`	Holds LLVM IR modules; only kernel libraries are turned into SPIR-V.
`VulkanCodegen`	Creates the LLVM modules and code libraries of the Vulkan target.

Functions#

`target_machine`()	The LLVM target machine for Vulkan SPIR-V, created on first use.
`emit_spirv`(llvm_ir[, exact, narrow_ints])	Translate LLVM IR to a SPIR-V binary in a child process.
`save_failure`(llvm_ir[, spirv])	Keep the input of a failed code generation for reproducing it.
`push_constant_layout`(spirv)	Offsets and sizes of the members of a module's push-constant block.
`spirv_capabilities`(spirv)	Find the optional capabilities a SPIR-V module declares.
`fix_compare_exchange`(spirv)	Repair the result of compare-and-swap operations.
`half_storage`(spirv)	Declare 16-bit storage instead of 16-bit arithmetic where possible.
`null_int_constants`(spirv)	Write integer zeros as `OpConstant` instead of `OpConstantNull`.
`storage8_capability`(spirv)	Declare `StorageBuffer8BitAccess` where buffers hold 8-bit integers.
`specialize_local_size`(spirv)	Make the workgroup size of a kernel a specialization constant.
`fix_barrier_semantics`(spirv)	Give barriers memory semantics that Vulkan accepts.
`strip_unused`(spirv)	Remove unused variables, constants and types from a SPIR-V module.
`narrow_index_constants`(spirv)	Turn 64-bit index constants in a SPIR-V module into 32-bit ones.
`mark_exact`(spirv)	Decorate all float arithmetic in a SPIR-V module as `NoContraction`.
`preserve_nan`(spirv, widths)	Require the driver to keep NaN, infinity and signed zero.
`lower_group_operations`(spirv)	Turn calls of subgroup placeholder functions into group instructions.
`check_structure`(spirv)	Check the rules of structured control flow that drivers rely on.
`check_spirv`(spirv)	Reject modules that are known to crash or confuse drivers.
`num_groups_members`(text, first)	Read the size of the grid from the push-constant block.
`local_size_array`(text, local)	Read the workgroup size from a private array instead of placeholders.

Module Contents#

numba_vulkan.codegen.TRIPLE = 'spirv1.5-unknown-vulkan1.2-compute'#

numba_vulkan.codegen.ENTRY_POINT = 'main'#

numba_vulkan.codegen.PASSES = ('always_inliner', 'global_dead_code_eliminate', 'sroa', 'instruction_combine', 'simplify_cfg',...#

numba_vulkan.codegen.EARLY_PASSES = ('always_inliner', 'global_dead_code_eliminate', 'sroa', 'instruction_combine', 'simplify_cfg',...#

numba_vulkan.codegen.EMIT_TIMEOUT = 120#

numba_vulkan.codegen.target_machine()#

The LLVM target machine for Vulkan SPIR-V, created on first use.

It is only used in this process to drive optimisation passes and to query the data layout; code generation happens in a child process.

Return type:: llvmlite.binding.TargetMachine

class numba_vulkan.codegen.Emitter(llc=None)#

A child process that runs LLVM’s SPIR-V backend.

The backend aborts the process on input it cannot handle, and it must not translate two modules in one process, so it runs outside the user’s interpreter. Starting a Python process that imports llvmlite takes a few tenths of a second; the child is therefore kept, and forks for every module (see numba_vulkan._emit). It is replaced when it dies or hangs. On platforms without fork a new child is started for every module.

With llc, LLVM’s llc program translates every module instead, in a process of its own. That serves when llvmlite’s own backend is broken, as it is in llvmlite 0.50’s macOS arm64 wheel (see docs/source/known_issues.md). Its LLVM must read the IR that llvmlite writes, so it should have the same major version.

Parameters:: llc (str, optional) – Path of llc; by default llvmlite’s backend is used.

llc = None#

warm_up()#

Start the child process in the background if it is not running.

Called when compilation of a function begins, so that the process is ready by the time there is a module to translate.

close()#: Terminate the child process; it is restarted when needed.

emit(llvm_ir)#

Translate a module.

Parameters:: llvm_ir (str) – Textual LLVM IR.
Returns:: The SPIR-V module.
Return type:: bytes
Raises:: SpirvCodegenError – If the backend fails or does not finish within EMIT_TIMEOUT seconds.

numba_vulkan.codegen.LLC_ENV_VAR = 'NUMBA_VULKAN_LLC'#

numba_vulkan.codegen.emitter#

numba_vulkan.codegen.emit_spirv(llvm_ir, exact=True, narrow_ints=False)#

Translate LLVM IR to a SPIR-V binary in a child process.

Parameters:

llvm_ir (str) – Textual LLVM IR of a module containing the shader entry point.
exact (bool) – Whether float arithmetic is marked exact; see mark_exact.
narrow_ints (bool) – Whether the module is meant for devices without 64-bit integers; see narrow_index_constants.

Returns:

The SPIR-V module.

Return type:

bytes

Raises:

SpirvCodegenError – If the backend fails or the module does not pass check_spirv. A note names the files the LLVM IR and the module were saved to, to reproduce the failure (see save_failure).

numba_vulkan.codegen.save_failure(llvm_ir, spirv=None)#

Keep the input of a failed code generation for reproducing it.

The files are named by a hash of the IR, in the failures folder of the cache directory (~/.cache/numba-vulkan), or of the system’s temporary directory if that cannot be written. llc -mtriple= with the triple in the IR, and spirv-val, reproduce the failure outside of Python.

Parameters:

llvm_ir (str) – The LLVM IR handed to the backend.
spirv (bytes, optional) – The module, if the backend produced one.

Returns:

The files written; empty if none could be.

Return type:

list of str

class numba_vulkan.codegen.CompiledKernel#

A kernel specialisation ready to be run by numba_vulkan.runtime.

Variables:

name (str) – Qualified name of the Python function.
spirv (bytes) – The SPIR-V module.
llvm_ir (str) – Optimised LLVM IR the module was generated from.
argtypes (tuple of numba.types.Type) – Argument types, with arrays bound to storage buffers.
ndim (int) – Dimensionality of the dispatch grid.
local_size (tuple of int) – Workgroup size along x, y and z.
capabilities (set of str) – Optional device features the shader needs, out of "float64", "int64", "int16" and "int8".
written_bindings (set of int) – Bindings whose buffers the shader stores to.
push_format (str) – struct format of the push-constant block; empty if the kernel has none and takes its scalars and extents from buffers.
push_sources (tuple of tuple) – What each value in the block is: ("arg", index) for a scalar argument, ("shape", k) for the k-th extent of all array arguments together, ("groups", axis) for the number of workgroups of the whole grid along an axis.
args_pushed (bool) – Whether scalars and extents are push constants rather than buffers.
exact (bool) – Whether the kernel was compiled without fastmath, so that floats follow IEEE rules, NaN included (see preserve_nan).

name: str#

spirv: bytes#

llvm_ir: str#

argtypes: tuple#

ndim: int#

local_size: tuple#

capabilities: set#

written_bindings: set#

read_bindings: set#

shared_bytes: int = 0#

print_binding: int = None#

constants: dict#

mode: numba_vulkan.narrowing.Mode#

narrowed: numba_vulkan.narrowing.Mode#

push_format: str = ''#

push_sources: tuple = ()#

args_pushed: bool = False#

exact: bool = True#

property push_size#

Size of the push-constant block in bytes; 0 without one.

Return type:: int

property unbound#

Bindings of arguments that are passed as push constants instead.

Return type:: set of int

property num_bindings#

Number of storage buffers the kernel uses.

Returns:: One for the shape buffer, one per argument, one per constant array, and one for the output of print.
Return type:: int

numba_vulkan.codegen.push_constant_layout(spirv)#

Offsets and sizes of the members of a module’s push-constant block.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: (offset, size) in bytes per member, in member order; empty if the module declares no push constants.
Return type:: list of tuple of int
Raises:: SpirvCodegenError – If a member is not a scalar or has no offset.

numba_vulkan.codegen.spirv_capabilities(spirv)#

Find the optional capabilities a SPIR-V module declares.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: Names matching the feature attributes of numba_vulkan.runtime.DeviceInfo, for example "float64".
Return type:: set of str

numba_vulkan.codegen.fix_compare_exchange(spirv)#

Repair the result of compare-and-swap operations.

LLVM’s llvm.spv.cmpxchg intrinsic yields the old value, but the backend packs it with a success flag into a composite that it then treats as the integer result. The composite is removed and its uses take the old value directly.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The repaired module; spirv itself if it has no compare-and-swap.
Return type:: bytes

numba_vulkan.codegen.half_storage(spirv)#

Declare 16-bit storage instead of 16-bit arithmetic where possible.

float16 arrays are read and written as half and computed with as float. The backend declares the Float16 capability for any use of half, which needs the shaderFloat16 device feature. If half values are only loaded, stored and converted, the capability is replaced by StorageBuffer16BitAccess, which more devices offer.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module; unchanged if it does not use half or computes with it.
Return type:: bytes

numba_vulkan.codegen.null_int_constants(spirv)#

Write integer zeros as OpConstant instead of OpConstantNull.

The backend emits OpConstantNull for an integer 0, also where it indexes into a structure, for which the SPIR-V specification asks for an OpConstant. Recent validators accept both; older ones, such as the one inside the validation layer of Ubuntu 24.04, reject the module.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module with every integer OpConstantNull replaced by an OpConstant of value 0 and the same id.
Return type:: bytes

numba_vulkan.codegen.storage8_capability(spirv)#

Declare StorageBuffer8BitAccess where buffers hold 8-bit integers.

Vulkan needs that capability, and the device feature of the same name, for 8-bit integers in storage buffers (int8 and uint8 arrays). The backend declares only Int8. Eight-bit values in other memory, such as the byte tables of libclc in private memory, do not need it.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module, with the capability added if a storage buffer contains an 8-bit integer.
Return type:: bytes

numba_vulkan.codegen.LOCAL_SIZE_ARRAY = 'nv.local_size'#

numba_vulkan.codegen.specialize_local_size(spirv)#

Make the workgroup size of a kernel a specialization constant.

One module then serves every workgroup size: the runtime sets specialization constants 0, 1 and 2 to its extents when it creates the pipeline. Their defaults are the size the module was compiled with. The size is declared as clspv declares it: three OpSpecConstant combined into an OpSpecConstantComposite decorated BuiltIn WorkgroupSize, which takes precedence over the LocalSize execution mode. Reads of the size in the kernel come from the private array LOCAL_SIZE_ARRAY, whose initializer becomes a composite of the same constants.

Parameters:: spirv (bytes) – A SPIR-V module with one LocalSize execution mode.
Returns:: The module with the constants added.
Return type:: bytes

numba_vulkan.codegen.fix_barrier_semantics(spirv)#

Give barriers memory semantics that Vulkan accepts.

The backend emits OpControlBarrier with sequentially consistent semantics, which Vulkan forbids. They become acquire-release on workgroup and buffer memory, which is what syncthreads promises.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The repaired module; spirv itself if it has no barrier.
Return type:: bytes

numba_vulkan.codegen.strip_unused(spirv)#

Remove unused variables, constants and types from a SPIR-V module.

LLVM’s backend wants a name string for every buffer, and emits each as a local array of 8-bit integers that nothing reads. Those arrays are the only reason most shaders declare the Int8 capability, an optional device feature. This removes local variables, constants and types that are never used, and then the capabilities for integer and float widths that no longer occur.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module without the unused definitions.
Return type:: bytes

Notes

Operands of instructions that _id_operands does not know all count as uses, so a literal can keep a definition alive, but nothing that is used is ever removed.

numba_vulkan.codegen.narrow_index_constants(spirv)#

Turn 64-bit index constants in a SPIR-V module into 32-bit ones.

The backend indexes into constant tables with 64-bit constants of its own making, whatever the width of the indices in the LLVM IR. In a kernel compiled for a device without 64-bit integers they are the last use of that type.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module with the constants retyped, or spirv itself if it uses 64-bit integers for anything but small constant indices of access chains.
Return type:: bytes

numba_vulkan.codegen.mark_exact(spirv)#

Decorate all float arithmetic in a SPIR-V module as NoContraction.

Without the decoration, shader compilers are free to reassociate float arithmetic: all tested drivers simplify (1 + x) - 1 to x, which destroys compensated algorithms such as the one behind log1p. Numba promises IEEE semantics unless fastmath is requested, so every add, subtract, multiply, divide, negate and remainder is marked exact.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module with the decorations inserted.
Return type:: bytes

numba_vulkan.codegen.preserve_nan(spirv, widths)#

Require the driver to keep NaN, infinity and signed zero.

Without the SignedZeroInfNanPreserve execution mode, a Vulkan driver may compile as if no float were NaN or infinite, and NVIDIA’s does: it turns y < x ? y : x (Python’s min) into an instruction that ignores NaN. Kernels compiled without fastmath therefore ask for these values to be kept, for the float widths they compute with.

Parameters:

spirv (bytes) – A SPIR-V module with one entry point.
widths (tuple of int) – The float widths the device can keep them for.

Returns:

The module with the capability and an execution mode per width it uses; spirv itself if there is none.

Return type:

bytes

numba_vulkan.codegen.GROUP_PREFIX = 'nv.sg.'#

numba_vulkan.codegen.lower_group_operations(spirv)#

Turn calls of subgroup placeholder functions into group instructions.

LLVM’s backend offers subgroup reductions but not products, scans or shuffles. Kernels therefore call functions named nv.sg.<operation>[.<red|incl|excl>].<type> that are only declared, which the backend emits as imported functions (with the Linkage capability, which Vulkan does not allow). Each call becomes the OpGroupNonUniform* instruction of that name, on the subgroup scope, and the declarations, their decorations and the capability go.

Parameters:: spirv (bytes) – A SPIR-V module.
Returns:: The module; spirv itself if it calls no such function.
Return type:: bytes
Raises:: SpirvCodegenError – If a placeholder has a name this pass does not know.

numba_vulkan.codegen.check_structure(spirv)#

Check the rules of structured control flow that drivers rely on.

A cheap subset of what spirv-val checks, applied to every module, so that a graph that neither numba_vulkan.structurize nor LLVM’s backend brought into shape is caught here rather than by a driver:

a conditional branch or switch is preceded by a merge instruction, unless it leaves or continues a loop;
a branch back to an earlier block goes to a loop header;
no block is the merge block of two headers.

Parameters:: spirv (bytes) – A SPIR-V module.
Raises:: SpirvCodegenError – If a rule is broken.

numba_vulkan.codegen.check_spirv(spirv)#

Reject modules that are known to crash or confuse drivers.

Parameters:: spirv (bytes) – A SPIR-V module.
Raises:: SpirvCodegenError – If the module selects between pointers (invalid without the VariablePointers capability), breaks a rule of structured control flow (see check_structure), or, with NUMBA_VULKAN_VALIDATE=1 in the environment, if spirv-val rejects it.

numba_vulkan.codegen.num_groups_members(text, first)#

Read the size of the grid from the push-constant block.

Parameters:

text (str) – Textual LLVM IR after optimisation, with numba_vulkan.num_groups placeholder calls.
first (int) – Number of the first member to use; the size takes three, after the arguments.

Returns:

text (str) – The IR, with the calls turned into placeholder loads from the block.
used (bool) – Whether the kernel reads the size of the grid.

numba_vulkan.codegen.local_size_array(text, local)#

Read the workgroup size from a private array instead of placeholders.

numba_vulkan.local_size calls stand for the workgroup size until optimisation is over, so that LLVM does not fold it into the code. They become loads from LOCAL_SIZE_ARRAY, whose initializer specialize_local_size turns into specialization constants.

Parameters:

text (str) – Textual LLVM IR after optimisation.
local (tuple of int) – The workgroup size the module is compiled with, the default of the constants.

Returns:

The rewritten IR; text itself if it does not read the size.

Return type:

str

class numba_vulkan.codegen.VulkanCodeLibrary(codegen, name)#

Holds LLVM IR modules; only kernel libraries are turned into SPIR-V.

One library is created per compiled function. The library of a kernel additionally contains the entry point and links the libraries of all functions the kernel calls.

Parameters:

codegen (VulkanCodegen) – The code generator that owns the library.
name (str) – Name of the library.

Variables:

written_bindings (set of int) – Bindings the linked code stores to; filled when the library is linked.

written_bindings#

read_bindings#

constants#

print_binding = None#

first_constant_binding = 0#

local_size = None#

push_types = []#

add_ir_module(module)#

Add a module to the library.

Parameters:: module (llvmlite.ir.Module or str) – An IR module, or textual LLVM IR.

add_linking_library(library)#

Make the code of another library available to this one.

Parameters:: library (VulkanCodeLibrary) – Library holding functions that this library calls.

finalize()#: Mark the library complete. Linking is deferred until SPIR-V is needed.

get_function(name)#

Look up a function defined or declared in this library.

Parameters:: name (str) – Mangled function name.
Return type:: llvmlite.ir.Function
Raises:: KeyError – If no module of the library contains the function.

get_llvm_str()#

Unoptimised LLVM IR of this library alone.

Return type:: str

get_optimized_llvm_str()#

LLVM IR of the whole kernel as handed to the SPIR-V backend.

Return type:: str

abstractmethod get_asm_str()#

Not available; disassemble get_spirv with spirv-dis instead.

Raises:: NotImplementedError – Always.

get_spirv(exact=True)#

SPIR-V binary of the kernel this library holds the entry point of.

Parameters:: exact (bool) – Whether float arithmetic is marked exact; see mark_exact.
Returns:: The SPIR-V module. The result is cached.
Return type:: bytes
Raises:: SpirvCodegenError – If code generation fails.

class numba_vulkan.codegen.VulkanCodegen(module_name)#

Creates the LLVM modules and code libraries of the Vulkan target.

Parameters:: module_name (str) – Name used for diagnostics.

magic_tuple()#

Values identifying this code generator for Numba’s caching.

Return type:: tuple