Python API Reference

class vkdispatch.AddressMode(*values)

Bases: Enum

Defines how to handle out-of-bounds addresses when accessing an image.

REPEAT

Repeat the image data when accessing out-of-bounds addresses.

Type:

int

MIRRORED_REPEAT

Mirror and repeat the image data when accessing out-of-bounds addresses.

Type:

int

CLAMP_TO_EDGE

Clamp out-of-bounds addresses to the edge of the image.

Type:

int

CLAMP_TO_BORDER

Clamp out-of-bounds addresses to a specific border color.

Type:

int

MIRROR_CLAMP_TO_EDGE

Mirror the image and clamp out-of-bounds addresses to the edge of the image.

Type:

int

CLAMP_TO_BORDER = 3
CLAMP_TO_EDGE = 2
MIRRORED_REPEAT = 1
MIRROR_CLAMP_TO_EDGE = 4
REPEAT = 0
exception vkdispatch.BackendUnavailableError(backend_name: str, message: str)

Bases: ImportError

class vkdispatch.BorderColor(*values)

Bases: Enum

Defines the border color used when clamping out-of-bounds addresses.

FLOAT_TRANSPARENT_BLACK

A fully transparent black border.

Type:

int

INT_TRANSPARENT_BLACK

A fully transparent black border.

Type:

int

FLOAT_OPAQUE_BLACK

An opaque black border with a specific alpha value.

Type:

int

INT_OPAQUE_BLACK

An opaque black border with a specific alpha value.

Type:

int

FLOAT_OPAQUE_WHITE

An opaque white border with a specific alpha value.

Type:

int

INT_OPAQUE_WHITE

An opaque white border with a specific alpha value.

Type:

int

FLOAT_OPAQUE_BLACK = 2
FLOAT_OPAQUE_WHITE = 4
FLOAT_TRANSPARENT_BLACK = 0
INT_OPAQUE_BLACK = 3
INT_OPAQUE_WHITE = 5
INT_TRANSPARENT_BLACK = 1
class vkdispatch.Buffer(shape: Tuple[int, ...], var_type: dtype, external_buffer: ExternalBufferInfo = None)

Bases: Handle, Generic[_ArgType]

Represents a contiguous block of memory on the GPU (or shared across multiple devices).

Buffers are the primary mechanism for transferring data between the host (CPU) and the device (GPU). They are typed using vkdispatch.dtype and support multi-dimensional shapes, similar to NumPy arrays.

Parameters:

  • shape (Tuple[int, ...]) – The dimensions of the buffer. Must be a tuple of 1, 2, or 3 integers.

  • var_type (vkdispatch.base.dtype.dtype) – The data type of the elements stored in the buffer.

Raises:

ValueError – If the shape has more than 3 dimensions or if the requested size exceeds 2^30 elements.

cuda_array_stream: Any | None
cuda_ptr: int | None
cuda_source: Any
is_external: bool
is_writable: bool
mem_size: int
owns_memory: bool
read(index: int | None = None)

Downloads data from the GPU buffer to the host.

Parameters:

index (Union[int, None]) – The device index to read from. If None, reads from all devices and returns a stacked array with an extra dimension for the device index.

Returns:

A host array representation containing the buffer data.

Raises:

ValueError – If the specified index is invalid.

shape: Tuple[int]
signals: List[Signal]
size: int
var_type: dtype
write(data: bytes | bytearray | memoryview | Any, index: int = None) → None

Uploads data from the host to the GPU buffer.

If index is None, the data is broadcast to the memory of all active devices in the context. Otherwise, it writes only to the device specified by the index.

Parameters:

  • data (Union[bytes, bytearray, memoryview, Any]) – The source data. Can be a bytes-like object or an array-like object.

  • index (int) – The device index to write to. Defaults to -1 (all devices).

Raises:

ValueError – If the data size exceeds the buffer size or if the index is invalid.

class vkdispatch.BufferBindInfo(buffer: Buffer, binding: int, shape_name: str, read_access: bool, write_access: bool)

Bases: object

A dataclass to hold information about a buffer binding.

binding: int
buffer: Buffer
read_access: bool
shape_name: str
write_access: bool
class vkdispatch.CUDAGraphCapture

Bases: object

add_uniform_buffer(buffer)
cuda_stream

alias of Any

uniform_buffers

alias of List[Any]

class vkdispatch.CommandGraph(reset_on_submit: bool = False, submit_on_record: bool = False)

Bases: CommandList

A high-level abstraction over CommandList that manages resource binding and push constants automatically.

Unlike a raw CommandList, a CommandGraph tracks variable state and handles the complexities of BufferBuilder for push constants and uniform buffers. It serves as the default recording target for shader functions.

Parameters:

  • reset_on_submit (bool) – If True, the graph clears its recorded commands immediately after submission.

  • submit_on_record (bool) – If True, commands are submitted to the GPU immediately upon recording (simulating immediate mode execution).

bind_var(name: str)
buffers_valid: bool
ensure_uniform_constants_capacity(uniform_word_size: int) → None
name_to_pc_key_dict: Dict[str, List[Tuple[str, str]]]
pc_builder: BufferBuilder
pc_values: Dict[Tuple[str, str], Any]
prepare_for_cuda_graph_capture(instance_count: int = None) → None

Initialize internal data uploads before torch CUDA graph capture.

This method performs one-time uniform/push-constant staging without submitting the command list, so only kernel launches are captured by torch.cuda.graph.

prepare_submission_state(instance_count: int) → None
queued_pc_values: Dict[Tuple[str, str], Any]
record_shader(plan: ComputePlan, shader_description: ShaderDescription, exec_limits: Tuple[int, int, int], blocks: Tuple[int, int, int], bound_buffers: List[BufferBindInfo], bound_samplers: List[ImageBindInfo], uniform_values: Dict[str, Any] = {}, pc_values: Dict[str, Any] = {}, shader_uuid: str = None) → None

Internal method to record a high-level shader execution.

This method handles the creation of DescriptorSet objects, binding of buffers and images, and populating push constant/uniform data before calling the base record_compute_plan.

Parameters:

  • plan – The compute plan to execute.

  • shader_description – Metadata about the shader source and layout.

  • exec_limits – The execution limits (grid size) in x, y, z.

  • blocks – The number of workgroups to dispatch.

  • bound_buffers – List of buffers to bind.

  • bound_samplers – List of images/samplers to bind.

  • uniform_values – Dictionary of values for uniform buffer objects.

  • pc_values – Dictionary of values for push constants.

  • shader_uuid – Unique identifier for this shader instance (for caching).

recorded_descriptor_sets: List[DescriptorSet]
reset() → None

Reset the command graph by clearing the push constant buffer and descriptor set lists.

set_var(name: str, value: Any)
submit(instance_count: int = None, queue_index: int = -2) → None

Submit the command list to the specified device with additional data to append to the front of the command list.

Parameters: device_index (int): The device index to submit the command list to.

Default is 0.

data (bytes): The additional data to append to the front of the command list.

submit_any(instance_count: int = None) → None
submit_on_record: bool
uniform_bindings: Any
uniform_builder: BufferBuilder
uniform_constants_buffer: Buffer | None
uniform_constants_size: int
uniform_descriptors: List[Tuple[DescriptorSet, int, int]]
uniform_values: Dict[Tuple[str, str], Any]
class vkdispatch.DeviceInfo(dev_index: int, version_variant: int, version_major: int, version_minor: int, version_patch: int, driver_version: int, vendor_id: int, device_id: int, device_type: int, device_name: str, shader_buffer_float32_atomics: int, shader_buffer_float32_atomic_add: int, float_64_support: int, float_16_support: int, int_64_support: int, int_16_support: int, storage_buffer_16_bit_access: int, uniform_and_storage_buffer_16_bit_access: int, storage_push_constant_16: int, storage_input_output_16: int, max_workgroup_size: Tuple[int, int, int], max_workgroup_invocations: int, max_workgroup_count: Tuple[int, int, int], max_bound_descriptor_sets: int, max_push_constant_size: int, max_storage_buffer_range: int, max_uniform_buffer_range: int, uniform_buffer_alignment: int, sub_group_size: int, supported_stages: int, supported_operations: int, quad_operations_in_all_stages: int, max_compute_shared_memory_size: int, queue_properties: List[Tuple[int, int]], scalar_block_layout: int, timeline_semaphores: int, uuid: bytes | None)

Bases: object

A class which represents all the features and properties of a Vulkan device.

NOTE: This class is not meant to be instantiated by the user. Instead, the user should call the get_devices() function to get a list of DeviceInfo instances.

dev_index

The index of the device.

Type:

int

version_variant

The Vulkan variant version.

Type:

int

version_major

The Vulkan major version.

Type:

int

version_minor

The Vulkan minor version.

Type:

int

version_patch

The Vulkan patch version.

Type:

int

driver_version

The version of the driver.

Type:

int

vendor_id

The vendor ID of the device.

Type:

int

device_id

The device ID of the device.

Type:

int

device_type

The device type, which is one of the following: 0: Other 1: Integrated GPU 2: Discrete GPU 3: Virtual GPU 4: CPU

Type:

int

device_name

The name of the device.

Type:

str

shader_buffer_float32_atomics

float32 atomics support.

Type:

int

shader_buffer_float32_atomic_add

float32 atomic add support.

Type:

int

float_64_support

64-bit float support.

Type:

int

int_64_support

64-bit integer support.

Type:

int

int_16_support

16-bit integer support.

Type:

int

max_workgroup_size

Maximum workgroup size.

Type:

Tuple[int, int, int]

max_workgroup_invocations

Maximum workgroup invocations.

Type:

int

max_workgroup_count

Maximum workgroup count.

Type:

Tuple[int, int, int]

max_bound_descriptor_sets

Maximum bound descriptor sets.

Type:

int

max_push_constant_size

Maximum push constant size.

Type:

int

max_storage_buffer_range

Maximum storage buffer range.

Type:

int

max_uniform_buffer_range

Maximum uniform buffer range.

Type:

int

uniform_buffer_alignment

Uniform buffer alignment.

Type:

int

sub_group_size

Subgroup size.

Type:

int

supported_stages

Supported subgroup stages.

Type:

int

supported_operations

Supported subgroup operations.

Type:

int

quad_operations_in_all_stages

Quad operations in all stages.

Type:

int

max_compute_shared_memory_size

Maximum compute shared memory size.

Type:

int

queue_properties

Queue properties.

Type:

List[Tuple[int, int]]

get_info_string(verbose: bool = False) → str

A method which returns a string representation of the device information.

Parameters:

verbose (bool) – A flag to enable verbose output.

Returns:

A string representation of the device information.

Return type:

str

is_apple() → bool

A method which checks if the device is an Apple device.

Returns:

A flag indicating whether the device is an Apple device.

Return type:

bool

is_nvidia() → bool

A method which checks if the device is an NVIDIA device.

Returns:

A flag indicating whether the device is an NVIDIA device.

Return type:

bool

class vkdispatch.Filter(*values)

Bases: Enum

Defines the filter used for image sampling.

NEAREST

Nearest neighbor filtering.

Type:

int

LINEAR

Linear interpolation filtering.

Type:

int

LINEAR = 1
NEAREST = 0
class vkdispatch.Image(shape: Tuple[int, ...], layers: int, dtype: type, channels: int, view_type: image_view_type, enable_mipmaps: bool = False)

Bases: Handle

read(device_index: int = 0)
sample(mag_filter: Filter = Filter.LINEAR, min_filter: Filter = Filter.LINEAR, mip_filter: Filter = Filter.LINEAR, address_mode: AddressMode = AddressMode.CLAMP_TO_EDGE, mip_lod_bias: float = 0.0, min_lod: float = 0.0, max_lod: float = 0.0, border_color: BorderColor = BorderColor.FLOAT_OPAQUE_WHITE) → Sampler
write(data: Any, device_index: int = -1) → None
class vkdispatch.Image1D(shape: int, dtype: type, channels: int = 1, enable_mipmaps: bool = False)

Bases: Image

class vkdispatch.Image2D(shape: ~typing.Tuple[int, int], dtype: type = <class 'vkdispatch.base.dtype._F32'>, channels: int = 1, enable_mipmaps: bool = False)

Bases: Image

class vkdispatch.Image2DArray(shape: ~typing.Tuple[int, int], layers: int, dtype: type = <class 'vkdispatch.base.dtype._F32'>, channels: int = 1, enable_mipmaps: bool = False)

Bases: Image

class vkdispatch.Image3D(shape: ~typing.Tuple[int, int, int], dtype: type = <class 'vkdispatch.base.dtype._F32'>, channels: int = 1, enable_mipmaps: bool = False)

Bases: Image

class vkdispatch.ImageBindInfo(sampler: Sampler, binding: int, read_access: bool, write_access: bool)

Bases: object

A dataclass to hold information about an image binding.

binding: int
read_access: bool
sampler: Sampler
write_access: bool
class vkdispatch.LogLevel(*values)

Bases: Enum

An enumeration which represents the log levels.

VERBOSE

All possible logs are printed. Module must be compiled with debug mode enabled to see these logs.

Type:

int

INFO

All release mode logs are printed. Useful for debugging the publicly available module.

Type:

int

WARNING

Only warnings and errors are printed. Default log level.

Type:

int

ERROR

Only errors are printed. Useful for muting annoying warnings that you know are harmless.

Type:

int

ERROR = 3
INFO = 1
VERBOSE = 0
WARNING = 2
class vkdispatch.MappingFunction(buffer_types: List[vkdispatch.base.dtype.dtype], return_type: vkdispatch.base.dtype.dtype, mapping_function: Callable, instance_id: uuid.UUID = <factory>)

Bases: object

buffer_types: List[dtype]
callback(*args)
instance_id: UUID
mapping_function: Callable
return_type: dtype
class vkdispatch.RFFTBuffer(shape: ~typing.Tuple[int, ...], fourier_type: ~vkdispatch.base.dtype.dtype = <class 'vkdispatch.base.dtype._CF64'>)

Bases: Buffer

fourier_shape: Tuple[int, ...]
read_fourier(index: int | None = None)
read_real(index: int | None = None)
real_shape: Tuple[int, ...]
real_type: dtype
write_fourier(data, index: int = None)
write_real(data, index: int = None)
class vkdispatch.Sampler(image: Image, mag_filter: Filter = Filter.LINEAR, min_filter: Filter = Filter.LINEAR, mip_filter: Filter = Filter.LINEAR, address_mode: AddressMode = AddressMode.CLAMP_TO_EDGE, mip_lod_bias: float = 0.0, min_lod: float = 0.0, max_lod: float = 0.0, border_color: BorderColor = BorderColor.FLOAT_OPAQUE_WHITE)

Bases: Handle

image: Image
exception vkdispatch.ShaderBuildError(message: str, *, shader_source: ShaderSource, compiler_log: str | None = None)

Bases: RuntimeError

compiler_log: str | None
shader_source: ShaderSource
class vkdispatch.ShaderFunction(shader_description: ShaderDescription, local_size=None, workgroups=None, exec_count=None)

Bases: object

bounds: ExecutionBounds
build()
exec_size: Tuple[int, int, int] | Callable | None
get_src(line_numbers: bool = None, build: bool = True) → ShaderSource
local_size: Tuple[int, int, int] | Callable | None
name: str
plan: ComputePlan
print_src(line_numbers: bool = None)
ready: bool
shader_description: ShaderDescription
source: str
workgroups: Tuple[int, int, int] | Callable | None
class vkdispatch.ShaderSource(name: str, code: str, local_size: Tuple[int, int, int])

Bases: object

code: str
local_size: Tuple[int, int, int]
name: str
class vkdispatch.Signal(ptr_addr: int = None)

Bases: object

free()
ptr_addr: int
try_wait(wait_for_timestamp: bool, queue_index: int)
wait(wait_for_timestamp: bool, queue_index: int)
vkdispatch.asbuffer(array: Any) → Buffer

Cast an array-like object to a buffer object.

vkdispatch.asrfftbuffer(data, fourier_type: dtype | None = None) → RFFTBuffer
vkdispatch.buffer_c64(shape: Tuple[int, ...]) → Buffer

Create a buffer of 64-bit complex numbers with the specified shape.

vkdispatch.buffer_dv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of 64-bit floating-point vectors of size 2 with the specified shape.

vkdispatch.buffer_dv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of 64-bit floating-point vectors of size 3 with the specified shape.

vkdispatch.buffer_dv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of 64-bit floating-point vectors of size 4 with the specified shape.

vkdispatch.buffer_f16(shape: Tuple[int, ...]) → Buffer

Create a buffer of 16-bit floating-point numbers with the specified shape.

vkdispatch.buffer_f32(shape: Tuple[int, ...]) → Buffer

Create a buffer of 32-bit floating-point numbers with the specified shape.

vkdispatch.buffer_f64(shape: Tuple[int, ...]) → Buffer

Create a buffer of 64-bit floating-point numbers with the specified shape.

vkdispatch.buffer_hv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of 16-bit floating-point vectors of size 2 with the specified shape.

vkdispatch.buffer_hv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of 16-bit floating-point vectors of size 3 with the specified shape.

vkdispatch.buffer_hv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of 16-bit floating-point vectors of size 4 with the specified shape.

vkdispatch.buffer_i16(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 16-bit integers with the specified shape.

vkdispatch.buffer_i32(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 32-bit integers with the specified shape.

vkdispatch.buffer_ihv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 16-bit integer vectors of size 2 with the specified shape.

vkdispatch.buffer_ihv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 16-bit integer vectors of size 3 with the specified shape.

vkdispatch.buffer_ihv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 16-bit integer vectors of size 4 with the specified shape.

vkdispatch.buffer_iv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 32-bit integer vectors of size 2 with the specified shape.

vkdispatch.buffer_iv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 32-bit integer vectors of size 3 with the specified shape.

vkdispatch.buffer_iv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of signed 32-bit integer vectors of size 4 with the specified shape.

vkdispatch.buffer_u16(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 16-bit integers with the specified shape.

vkdispatch.buffer_u32(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 32-bit integers with the specified shape.

vkdispatch.buffer_uhv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 16-bit integer vectors of size 2 with the specified shape.

vkdispatch.buffer_uhv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 16-bit integer vectors of size 3 with the specified shape.

vkdispatch.buffer_uhv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 16-bit integer vectors of size 4 with the specified shape.

vkdispatch.buffer_uv2(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 32-bit integer vectors of size 2 with the specified shape.

vkdispatch.buffer_uv3(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 32-bit integer vectors of size 3 with the specified shape.

vkdispatch.buffer_uv4(shape: Tuple[int, ...]) → Buffer

Create a buffer of unsigned 32-bit integer vectors of size 4 with the specified shape.

vkdispatch.buffer_v2(shape: Tuple[int, ...]) → Buffer

Create a buffer of 32-bit floating-point vectors of size 2 with the specified shape.

vkdispatch.buffer_v3(shape: Tuple[int, ...]) → Buffer

Create a buffer of 32-bit floating-point vectors of size 3 with the specified shape.

vkdispatch.buffer_v4(shape: Tuple[int, ...]) → Buffer

Create a buffer of 32-bit floating-point vectors of size 4 with the specified shape.

vkdispatch.complex128

alias of _CF128

vkdispatch.complex32

alias of _CF32

vkdispatch.complex64

alias of _CF64

vkdispatch.cuda_graph_capture(cuda_stream=None)
vkdispatch.default_graph() → CommandGraph
class vkdispatch.dtype

Bases: object

child_count: int
child_type: dtype
dimentions: int
format_str: str
glsl_type: str
item_size: int
name: str
numpy_shape: tuple
scalar: dtype | None
shape: tuple
true_numpy_shape: tuple
vkdispatch.dvec2

alias of _V2F64

vkdispatch.dvec3

alias of _V3F64

vkdispatch.dvec4

alias of _V4F64

vkdispatch.float16

alias of _F16

vkdispatch.float32

alias of _F32

vkdispatch.float64

alias of _F64

vkdispatch.from_cuda_array(obj: Any, var_type: dtype | None = None, require_contiguous: bool = True, writable: bool | None = None, keepalive: bool = True) → Buffer
vkdispatch.get_backend() → str
vkdispatch.get_context() → Context
vkdispatch.get_context_handle() → int
vkdispatch.get_cuda_capture() → CUDAGraphCapture
vkdispatch.get_devices() → List[DeviceInfo]

Get a list of DeviceInfo instances representing all the Vulkan devices on the system.

Returns:

A list of DeviceInfo instances.

Return type:

List[DeviceInfo]

vkdispatch.global_graph() → CommandGraph
vkdispatch.hvec2

alias of _V2F16

vkdispatch.hvec3

alias of _V3F16

vkdispatch.hvec4

alias of _V4F16

vkdispatch.ihvec2

alias of _V2I16

vkdispatch.ihvec3

alias of _V3I16

vkdispatch.ihvec4

alias of _V4I16

class vkdispatch.image_format(*values)

Bases: Enum

R16G16B16A16_SFLOAT = 97
R16G16B16A16_SINT = 96
R16G16B16A16_UINT = 95
R16G16B16_SFLOAT = 90
R16G16B16_SINT = 89
R16G16B16_UINT = 88
R16G16_SFLOAT = 83
R16G16_SINT = 82
R16G16_UINT = 81
R16_SFLOAT = 76
R16_SINT = 75
R16_UINT = 74
R32G32B32A32_SFLOAT = 109
R32G32B32A32_SINT = 108
R32G32B32A32_UINT = 107
R32G32B32_SFLOAT = 106
R32G32B32_SINT = 105
R32G32B32_UINT = 104
R32G32_SFLOAT = 103
R32G32_SINT = 102
R32G32_UINT = 101
R32_SFLOAT = 100
R32_SINT = 99
R32_UINT = 98
R64G64B64A64_SFLOAT = 121
R64G64B64A64_SINT = 120
R64G64B64A64_UINT = 119
R64G64B64_SFLOAT = 118
R64G64B64_SINT = 117
R64G64B64_UINT = 116
R64G64_SFLOAT = 115
R64G64_SINT = 114
R64G64_UINT = 113
R64_SFLOAT = 112
R64_SINT = 111
R64_UINT = 110
R8G8B8A8_SINT = 42
R8G8B8A8_UINT = 41
R8G8B8_SINT = 28
R8G8B8_UINT = 27
R8G8_SINT = 21
R8G8_UINT = 20
R8_SINT = 14
R8_UINT = 13
class vkdispatch.image_type(*values)

Bases: Enum

Defines the type of an image.

TYPE_1D

A 1-dimensional image.

Type:

int

TYPE_2D

A 2-dimensional image.

Type:

int

TYPE_3D

A 3-dimensional image.

Type:

int

TYPE_1D = (0,)
TYPE_2D = (1,)
TYPE_3D = (2,)
class vkdispatch.image_view_type(*values)

Bases: Enum

Defines the type of an image view.

VIEW_TYPE_1D

A 1-dimensional image view.

Type:

int

VIEW_TYPE_2D

A 2-dimensional image view.

Type:

int

VIEW_TYPE_3D

A 3-dimensional image view.

Type:

int

VIEW_TYPE_2D_ARRAY

A 2D array of images.

Type:

int

VIEW_TYPE_1D = (0,)
VIEW_TYPE_2D = (1,)
VIEW_TYPE_2D_ARRAY = (5,)
VIEW_TYPE_3D = (2,)
vkdispatch.initialize(debug_mode: bool = False, log_level: LogLevel = LogLevel.WARNING, loader_debug_logs: bool = False, backend: str | None = None)

A function which initializes the Vulkan dispatch library.

Parameters:

  • debug_mode (bool) – A flag to enable debug mode.

  • log_level (LogLevel) – The log level, which is one of the following: LogLevel.VERBOSE LogLevel.INFO LogLevel.WARNING LogLevel.ERROR

  • loader_debug_logs (bool) – A flag to enable vulkan loader debug logs.

  • backend (Optional[str]) – Runtime backend to use. Supported values are “vulkan”, “cuda”, “opencl”, and “dummy”. If omitted, the currently selected backend is reused. If no backend was selected yet, VKDISPATCH_BACKEND is used when set, otherwise “vulkan” is used.

vkdispatch.int16

alias of _I16

vkdispatch.int32

alias of _I32

vkdispatch.int64

alias of _I64

vkdispatch.is_context_initialized() → bool
vkdispatch.is_cuda() → bool

A function which checks if the active backend is a CUDA backend.

Returns:

A flag indicating whether the active backend is a CUDA backend.

Return type:

bool

vkdispatch.is_dummy() → bool

A function which checks if the active backend is the dummy backend.

Returns:

A flag indicating whether the active backend is the dummy backend.

Return type:

bool

vkdispatch.is_initialized() → bool

A function which checks if the Vulkan dispatch library has been initialized.

Returns:

A flag indicating whether the Vulkan dispatch library has been initialized.

Return type:

bool

vkdispatch.is_opencl() → bool

A function which checks if the active backend is the OpenCL backend.

Returns:

A flag indicating whether the active backend is the OpenCL backend.

Return type:

bool

vkdispatch.is_vulkan() → bool

A function which checks if the active backend is the Vulkan backend.

Returns:

A flag indicating whether the active backend is the Vulkan backend.

Return type:

bool

vkdispatch.ivec2

alias of _V2I32

vkdispatch.ivec3

alias of _V3I32

vkdispatch.ivec4

alias of _V4I32

vkdispatch.log(text: str, end: str = '\n', level: LogLevel = LogLevel.ERROR, stack_offset: int = 1)

A function which logs a message at the specified log level.

Parameters:

  • level (LogLevel) – The log level.

  • message (str) – The message to log.

vkdispatch.log_error(text: str, end: str = '\n')

A function which logs an error message.

Parameters:

message (str) – The message to log.

vkdispatch.log_info(text: str, end: str = '\n')

A function which logs an info message.

Parameters:

message (str) – The message to log.

vkdispatch.log_verbose(text: str, end: str = '\n')

A function which logs a verbose message.

Parameters:

message (str) – The message to log.

vkdispatch.log_warning(text: str, end: str = '\n')

A function which logs a warning message.

Parameters:

message (str) – The message to log.

vkdispatch.make_context(device_ids: int | List[int] | None = None, device_count: int | None = None, queue_counts: int | List[int] | None = None, queue_families: List[List[int]] | None = None, use_cpu: bool = False, multi_device: bool = False, multi_queue: bool = False) → Context
vkdispatch.make_shader_function(description: ShaderDescription, local_size=None, workgroups=None, exec_count=None) → ShaderFunction
vkdispatch.map(func: Callable, return_type: dtype = None, input_types: List[dtype] = None) → MappingFunction
vkdispatch.mat2

alias of _M2F32

vkdispatch.mat3

alias of _M3F32

vkdispatch.mat4

alias of _M4F32

vkdispatch.queue_wait_idle(queue_index: int = None, context: Context = None) → None

Wait for the specified queue to finish processing. For all queues, leave queue_index as None.

Parameters:

queue_index (int) – The index of the queue.

vkdispatch.select_queue_families(device_index: int, queue_count: int = None) → List[int]
vkdispatch.set_dummy_context_params(subgroup_size: int = None, subgroup_enabled: bool = None, max_workgroup_size: Tuple[int, int, int] = None, max_workgroup_invocations: int = None, max_workgroup_count: Tuple[int, int, int] = None, max_shared_memory: int = None) → None

Update cached context/device limit values for the active dummy backend context.

This only works when a dummy context already exists.

vkdispatch.set_global_graph(graph: CommandGraph = None) → CommandGraph
vkdispatch.set_log_level(level: LogLevel)

A function which sets the log level.

Parameters:

level (LogLevel) – The log level.

vkdispatch.shader(exec_size=None, local_size=None, workgroups=None, flags: ~vkdispatch.codegen.builder.ShaderFlags = <ShaderFlags.NONE: 0>, arg_type_annotations: ~typing.List[~typing.Any] | None = None, arg_names: ~typing.List[str] | None = None, arg_defaults: ~typing.List[~typing.Any] | None = None)

A decorator that transforms a Python function into a GPU Compute Shader.

The decorated function will undergo runtime inspection. Operations performed on vkdispatch types (buffers, registers) within the function are recorded and transpiled to GLSL.

Parameters:

  • exec_size (Union[int, Tuple[int, ...], Callable]) – The total number of threads to dispatch (x, y, z). The number of workgroups is calculated automatically based on local_size. Mutually exclusive with workgroups.

  • local_size (Union[int, Tuple[int, ...]]) – The number of threads per workgroup (x, y, z). Defaults to the device’s maximum supported workgroup size.

  • workgroups (Union[int, Tuple[int, ...], Callable]) – The explicit number of workgroups to dispatch (x, y, z). Mutually exclusive with exec_size.

  • flags (vkdispatch.codegen.ShaderFlags) – Compilation flags (e.g., vc.ShaderFlags.NO_EXEC_BOUNDS).

  • arg_type_annotations (Optional[List[Any]]) – Optional list of type annotations for the shader function’s parameters. If not provided, annotations will be inferred from the decorated function’s signature.

  • arg_names (Optional[List[str]]) – Optional list of parameter names corresponding to the type annotations. If not provided, names will be inferred from the decorated function’s signature.

  • arg_defaults (Optional[List[Any]]) – Optional list of default values for the parameters. If not provided, defaults will be inferred from the decorated function’s signature.

Returns:

A ShaderFunction wrapper that can be called to execute the kernel.

Raises:

ValueError – If both exec_size and workgroups are provided.

vkdispatch.uhvec2

alias of _V2U16

vkdispatch.uhvec3

alias of _V3U16

vkdispatch.uhvec4

alias of _V4U16

vkdispatch.uint16

alias of _U16

vkdispatch.uint32

alias of _U32

vkdispatch.uint64

alias of _U64

vkdispatch.uvec2

alias of _V2U32

vkdispatch.uvec3

alias of _V3U32

vkdispatch.uvec4

alias of _V4U32

vkdispatch.vec2

alias of _V2F32

vkdispatch.vec3

alias of _V3F32

vkdispatch.vec4

alias of _V4F32

class vkdispatch.codegen.BindingType(*values)

Bases: Enum

A dataclass that represents the type of a binding in a shader. Either a STORAGE_BUFFER, UNIFORM_BUFFER, or SAMPLER.

SAMPLER = 5
STORAGE_BUFFER = 1
UNIFORM_BUFFER = 3
class vkdispatch.codegen.BoundVariable(var_type: dtype, binding: int, name: str)

Bases: ShaderVariable

binding: int = -1
vkdispatch.codegen.Buff

alias of Buffer

class vkdispatch.codegen.Buffer

Bases: BufferVariable, Generic[_ArgType]

class vkdispatch.codegen.BufferVariable(var_type: dtype, binding: int, name: str, shape_var: ShaderVariable = None, shape_var_factory: Callable[[], ShaderVariable] | None = None, shape_name: str | None = None, raw_name: str | None = None, scalar_expr: str | None = None, codegen_backend: object | None = None, read_lambda: Callable[[], None] = None, write_lambda: Callable[[], None] = None)

Bases: BoundVariable

codegen_backend: object | None
read_callback()
read_lambda: Callable[[], None]
scalar_expr: str | None
property shape: ShaderVariable
write_callback()
write_lambda: Callable[[], None]
class vkdispatch.codegen.CUDABackend

Bases: CodeGenBackend

atomic_add_expr(mem_expr: str, value_expr: str, var_type: dtype) → str
barrier_statement() → str
binary_math_expr(func_name: str, lhs_type: dtype, lhs_expr: str, rhs_type: dtype, rhs_expr: str) → str
component_access_expr(expr: str, component: str, base_type: dtype) → str
constant_namespace() → str
constructor(var_type: dtype, args: List[str], arg_types: List[dtype | None] | None = None) → str
entry_point(body_contents: str) → str
exec_bounds_guard(exec_count_expr: str) → str
float_bits_to_int_expr(var_expr: str) → str
float_bits_to_uint_expr(var_expr: str) → str
fma_function_name(var_type: dtype) → str
global_invocation_id_expr() → str
group_memory_barrier_statement() → str
inf_f16_expr() → str
inf_f32_expr() → str
inf_f64_expr() → str
int_bits_to_float_expr(var_expr: str) → str
local_invocation_id_expr() → str
local_invocation_index_expr() → str
make_source(header: str, body: str, x: int, y: int, z: int) → str
mark_composite_binary_op(lhs_type: dtype, rhs_type: dtype, op: str, *, inplace: bool = False) → None
mark_composite_unary_op(var_type: dtype, op: str) → None
mark_feature_usage(feature_name: str) → None
math_func_name(func_name: str, var_type: dtype) → str

Return the backend-specific function name for a math operation.

Backends can override this to remap function names for specific types (e.g. CUDA __half intrinsics).

memory_barrier_buffer_statement() → str
memory_barrier_image_statement() → str
memory_barrier_shared_statement() → str
memory_barrier_statement() → str
name: str = 'cuda'
ninf_f16_expr() → str
ninf_f32_expr() → str
ninf_f64_expr() → str
num_subgroups_expr() → str
num_workgroups_expr() → str
pre_header(*, enable_subgroup_ops: bool, enable_printf: bool) → str
printf_statement(fmt: str, args: List[str]) → str
push_constant_declaration(contents: str) → str
reset_state() → None
sample_texture_expr(texture_expr: str, coord_expr: str, lod_expr: str | None = None) → str
sampler_declaration(binding: int, dimensions: int, name: str) → str
shared_buffer_declaration(var_type: dtype, name: str, size: int) → str
storage_buffer_declaration(binding: int, var_type: dtype, name: str) → str
subgroup_add_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_and_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_barrier_statement() → str
subgroup_elect_expr() → str
subgroup_id_expr() → str
subgroup_invocation_id_expr() → str
subgroup_max_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_min_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_mul_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_or_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_size_expr() → str
subgroup_xor_expr(arg_expr: str, arg_type: dtype | None = None) → str
texture_size_expr(texture_expr: str, lod: int, dimensions: int) → str
type_name(var_type: dtype) → str
uint_bits_to_float_expr(var_expr: str) → str
unary_math_expr(func_name: str, arg_type: dtype, arg_expr: str) → str
uniform_block_declaration(contents: str) → str
uses_feature(feature_name: str) → bool
variable_namespace() → str
workgroup_id_expr() → str
workgroup_size_expr() → str
class vkdispatch.codegen.CodeGenBackend

Bases: object

Interface for backend-specific code generation.

Subclasses should override all methods that are used by the codegen pipeline. The base implementation raises NotImplementedError so placeholder backends can be defined incrementally.

arithmetic_binary_expr(op: str, lhs_type: dtype, lhs_expr: str, rhs_type: dtype, rhs_expr: str) → str | None

Optional backend override for binary arithmetic expressions.

arithmetic_unary_expr(op: str, var_type: dtype, var_expr: str) → str | None

Optional backend override for unary arithmetic expressions.

atomic_add_expr(mem_expr: str, value_expr: str, var_type: dtype) → str
barrier_statement() → str
binary_math_expr(func_name: str, lhs_type: dtype, lhs_expr: str, rhs_type: dtype, rhs_expr: str) → str
buffer_component_expr(scalar_buffer_expr: str, base_type: dtype, element_index_expr: str, component_index_expr: str) → str | None
component_access_expr(expr: str, component: str, base_type: dtype) → str
constant_namespace() → str
constructor(var_type: dtype, args: List[str], arg_types: List[dtype | None] | None = None) → str
entry_point(body_contents: str) → str
exec_bounds_guard(exec_count_expr: str) → str
float_bits_to_int_expr(var_expr: str) → str
float_bits_to_uint_expr(var_expr: str) → str
fma_function_name(var_type: dtype) → str
global_invocation_id_expr() → str
group_memory_barrier_statement() → str
inf_f16_expr() → str
inf_f32_expr() → str
inf_f64_expr() → str
int_bits_to_float_expr(var_expr: str) → str
local_invocation_id_expr() → str
local_invocation_index_expr() → str
make_source(header: str, body: str, x: int, y: int, z: int) → str
mark_composite_binary_op(lhs_type: dtype, rhs_type: dtype, op: str, *, inplace: bool = False) → None
mark_composite_unary_op(var_type: dtype, op: str) → None
mark_feature_usage(feature_name: str) → None
mark_texture_sample_dimension(dimensions: int) → None
math_func_name(func_name: str, var_type: dtype) → str

Return the backend-specific function name for a math operation.

Backends can override this to remap function names for specific types (e.g. CUDA __half intrinsics).

memory_barrier_buffer_statement() → str
memory_barrier_image_statement() → str
memory_barrier_shared_statement() → str
memory_barrier_statement() → str
name: str = 'base'
ninf_f16_expr() → str
ninf_f32_expr() → str
ninf_f64_expr() → str
num_subgroups_expr() → str
num_workgroups_expr() → str
packed_buffer_read_expr(scalar_buffer_expr: str, var_type: dtype, element_index_expr: str) → str | None
packed_buffer_write_statements(scalar_buffer_expr: str, var_type: dtype, element_index_expr: str, value_expr: str) → str | None
pre_header(*, enable_subgroup_ops: bool, enable_printf: bool) → str
printf_statement(fmt: str, args: List[str]) → str
push_constant_declaration(contents: str) → str
reset_state() → None
sample_texture_expr(texture_expr: str, coord_expr: str, lod_expr: str | None = None) → str
sampler_declaration(binding: int, dimensions: int, name: str) → str
shared_buffer_declaration(var_type: dtype, name: str, size: int) → str
storage_buffer_declaration(binding: int, var_type: dtype, name: str) → str
subgroup_add_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_and_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_barrier_statement() → str
subgroup_elect_expr() → str
subgroup_id_expr() → str
subgroup_invocation_id_expr() → str
subgroup_max_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_min_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_mul_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_or_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_size_expr() → str
subgroup_xor_expr(arg_expr: str, arg_type: dtype | None = None) → str
texture_size_expr(texture_expr: str, lod: int, dimensions: int) → str
type_name(var_type: dtype) → str
uint_bits_to_float_expr(var_expr: str) → str
unary_math_expr(func_name: str, arg_type: dtype, arg_expr: str) → str
uniform_block_declaration(contents: str) → str
uses_feature(feature_name: str) → bool
variable_namespace() → str
workgroup_id_expr() → str
workgroup_size_expr() → str
vkdispatch.codegen.Const

alias of Constant

class vkdispatch.codegen.Constant

Bases: ShaderVariable, Generic[_ArgType]

class vkdispatch.codegen.GLSLBackend

Bases: CodeGenBackend

atomic_add_expr(mem_expr: str, value_expr: str, var_type: dtype) → str
barrier_statement() → str
constant_namespace() → str
constructor(var_type: dtype, args: List[str], arg_types: List[dtype | None] | None = None) → str
entry_point(body_contents: str) → str
exec_bounds_guard(exec_count_expr: str) → str
float_bits_to_int_expr(var_expr: str) → str
float_bits_to_uint_expr(var_expr: str) → str
global_invocation_id_expr() → str
group_memory_barrier_statement() → str
inf_f16_expr() → str
inf_f32_expr() → str
inf_f64_expr() → str
int_bits_to_float_expr(var_expr: str) → str
local_invocation_id_expr() → str
local_invocation_index_expr() → str
make_source(header: str, body: str, x: int, y: int, z: int) → str
mark_feature_usage(feature_name: str) → None
memory_barrier_buffer_statement() → str
memory_barrier_image_statement() → str
memory_barrier_shared_statement() → str
memory_barrier_statement() → str
name: str = 'glsl'
ninf_f16_expr() → str
ninf_f32_expr() → str
ninf_f64_expr() → str
num_subgroups_expr() → str
num_workgroups_expr() → str
pre_header(*, enable_subgroup_ops: bool, enable_printf: bool) → str
printf_statement(fmt: str, args: List[str]) → str
push_constant_declaration(contents: str) → str
reset_state() → None
sample_texture_expr(texture_expr: str, coord_expr: str, lod_expr: str | None = None) → str
sampler_declaration(binding: int, dimensions: int, name: str) → str
shared_buffer_declaration(var_type: dtype, name: str, size: int) → str
storage_buffer_declaration(binding: int, var_type: dtype, name: str) → str
subgroup_add_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_and_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_barrier_statement() → str
subgroup_elect_expr() → str
subgroup_id_expr() → str
subgroup_invocation_id_expr() → str
subgroup_max_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_min_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_mul_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_or_expr(arg_expr: str, arg_type: dtype | None = None) → str
subgroup_size_expr() → str
subgroup_xor_expr(arg_expr: str, arg_type: dtype | None = None) → str
texture_size_expr(texture_expr: str, lod: int, dimensions: int) → str
type_name(var_type: dtype) → str
uint_bits_to_float_expr(var_expr: str) → str
uniform_block_declaration(contents: str) → str
uses_feature(feature_name: str) → bool
variable_namespace() → str
workgroup_id_expr() → str
workgroup_size_expr() → str
class vkdispatch.codegen.Image1D

Bases: ImageVariable, Generic[_ArgType]

dimensions: int = 1
class vkdispatch.codegen.Image2D

Bases: ImageVariable, Generic[_ArgType]

dimensions: int = 2
class vkdispatch.codegen.Image3D

Bases: ImageVariable, Generic[_ArgType]

dimensions: int = 3
class vkdispatch.codegen.ImageVariable(var_type: dtype, binding: int, dimensions: int, name: str, read_lambda: Callable[[], None] = None, write_lambda: Callable[[], None] = None)

Bases: BoundVariable

dimensions: int = 0
read_callback()
read_lambda: Callable[[], None]
sample(coord: ShaderVariable, lod: ShaderVariable = None) → ShaderVariable
write_callback()
write_lambda: Callable[[], None]
vkdispatch.codegen.Img1

alias of Image1D

vkdispatch.codegen.Img2

alias of Image2D

vkdispatch.codegen.Img3

alias of Image3D

class vkdispatch.codegen.OpenCLBackend

Bases: CodeGenBackend

arithmetic_binary_expr(op: str, lhs_type: dtype, lhs_expr: str, rhs_type: dtype, rhs_expr: str) → str | None

Optional backend override for binary arithmetic expressions.

arithmetic_unary_expr(op: str, var_type: dtype, var_expr: str) → str | None

Optional backend override for unary arithmetic expressions.

atomic_add_expr(mem_expr: str, value_expr: str, var_type: dtype) → str
barrier_statement() → str
binary_math_expr(func_name: str, lhs_type: dtype, lhs_expr: str, rhs_type: dtype, rhs_expr: str) → str
buffer_component_expr(scalar_buffer_expr: str, base_type: dtype, element_index_expr: str, component_index_expr: str) → str | None
component_access_expr(expr: str, component: str, base_type: dtype) → str
constant_namespace() → str
constructor(var_type: dtype, args: List[str], arg_types: List[dtype | None] | None = None) → str
entry_point(body_contents: str) → str
exec_bounds_guard(exec_count_expr: str) → str
float_bits_to_int_expr(var_expr: str) → str
float_bits_to_uint_expr(var_expr: str) → str
global_invocation_id_expr() → str
group_memory_barrier_statement() → str
inf_f16_expr() → str
inf_f32_expr() → str
inf_f64_expr() → str
int_bits_to_float_expr(var_expr: str) → str
local_invocation_id_expr() → str
local_invocation_index_expr() → str
make_source(header: str, body: str, x: int, y: int, z: int) → str
mark_texture_sample_dimension(dimensions: int) → None
math_func_name(func_name: str, var_type: dtype) → str

Return the backend-specific function name for a math operation.

Backends can override this to remap function names for specific types (e.g. CUDA __half intrinsics).

memory_barrier_buffer_statement() → str
memory_barrier_image_statement() → str
memory_barrier_shared_statement() → str
memory_barrier_statement() → str
name: str = 'opencl'
ninf_f16_expr() → str
ninf_f32_expr() → str
ninf_f64_expr() → str
num_subgroups_expr() → str
num_workgroups_expr() → str
packed_buffer_read_expr(scalar_buffer_expr: str, var_type: dtype, element_index_expr: str) → str | None
packed_buffer_write_statements(scalar_buffer_expr: str, var_type: dtype, element_index_expr: str, value_expr: str) → str | None
pre_header(*, enable_subgroup_ops: bool, enable_printf: bool) → str
printf_statement(fmt: str, args: List[str]) → str
push_constant_declaration(contents: str) → str
reset_state() → None
sample_texture_expr(texture_expr: str, coord_expr: str, lod_expr: str | None = None) → str
sampler_declaration(binding: int, dimensions: int, name: str) → str
shared_buffer_declaration(var_type: dtype, name: str, size: int) → str
storage_buffer_declaration(binding: int, var_type: dtype, name: str) → str
subgroup_add_expr(arg_expr: str) → str
subgroup_and_expr(arg_expr: str) → str
subgroup_barrier_statement() → str
subgroup_elect_expr() → str
subgroup_id_expr() → str
subgroup_invocation_id_expr() → str
subgroup_max_expr(arg_expr: str) → str
subgroup_min_expr(arg_expr: str) → str
subgroup_mul_expr(arg_expr: str) → str
subgroup_or_expr(arg_expr: str) → str
subgroup_size_expr() → str
subgroup_xor_expr(arg_expr: str) → str
texture_size_expr(texture_expr: str, lod: int, dimensions: int) → str
type_name(var_type: dtype) → str
uint_bits_to_float_expr(var_expr: str) → str
unary_math_expr(func_name: str, arg_type: dtype, arg_expr: str) → str
uniform_block_declaration(contents: str) → str
variable_namespace() → str
workgroup_id_expr() → str
workgroup_size_expr() → str
class vkdispatch.codegen.ShaderArgumentInfo(name: str, arg_type: vkdispatch.codegen.shader_description.ShaderArgumentType, default_value: Any, shader_name: str | Dict[str, str], shader_shape_name: str | None = None, binding: int | None = None)

Bases: object

arg_type: ShaderArgumentType
binding: int | None
default_value: Any
name: str
shader_name: str | Dict[str, str]
shader_shape_name: str | None
class vkdispatch.codegen.ShaderBinding(dtype: dtype, name: str, dimension: int, binding_type: BindingType)

Bases: object

A dataclass that represents a bound resource in a shader. Either a buffer or an image.

dtype

The dtype of the resource. If the resource is an image, this should be vd.vec4 (since all images are sampled with 4 channels in shaders).

Type:

vd.dtype

name

The name of the resource within the shader code.

Type:

str

dimension

The dimension of the resource. Set to 0 for buffers and 1, 2, or 3 for images.

Type:

int

binding_type

The type of the binding. Either STORAGE_BUFFER, UNIFORM_BUFFER, or SAMPLER.

Type:

BindingType

binding_type: BindingType
dimension: int
dtype: dtype
name: str
class vkdispatch.codegen.ShaderBuilder(flags: ~vkdispatch.codegen.builder.ShaderFlags = <ShaderFlags.NONE: 0>)

Bases: ShaderWriter

backend: CodeGenBackend
binding_count: int
binding_list: List[ShaderBinding]
binding_read_access: Dict[int, bool]
binding_write_access: Dict[int, bool]
build(name: str) → ShaderDescription
compose_struct_decleration(elements: List[StructElement]) → str
declare_shader_arguments(type_annotations: List, names: List[str] | None = None, defaults: List[Any] | None = None)
exec_count: ShaderVariable | None
flags: ShaderFlags
get_shader_arguments() → List[ShaderVariable]
has_ubo: bool
new_scaled_var(var_type: dtype, name: str, scale: int = 1, offset: int = 0, parents: List[BaseVariable] = None)
new_var(var_type: dtype, name: str, parents: List[ShaderVariable], lexical_unit: bool = False, settable: bool = False, register: bool = False) → ShaderVariable
pc_struct: StructBuilder
reset() → None
scope_num: int
shader_arg_infos: List[ShaderArgumentInfo]
shader_args: List[ShaderVariable]
shared_buffer(var_type: dtype, size: int, var_name: str | None = None)
shared_buffers: List[SharedBuffer]
uniform_struct: StructBuilder
class vkdispatch.codegen.ShaderContext(builder: ShaderBuilder)

Bases: object

builder: ShaderBuilder
declare_input_arguments(type_annotations: List, names: List[str] | None = None, defaults: List[Any] | None = None)
get_description(name: str | None = None)
shader_description: ShaderDescription
class vkdispatch.codegen.ShaderDescription(header: str, body: str, name: str, pc_size: int, pc_structure: List[StructElement], uniform_structure: List[StructElement], binding_type_list: List[BindingType], binding_access: List[Tuple[bool, bool]], exec_count_name: str | None, resource_binding_base: int, backend: CodeGenBackend, shader_arg_infos: List[ShaderArgumentInfo])

Bases: object

A dataclass that represents a description of a shader object.

source

The source code of the shader.

Type:

str

pc_size

The size of the push constant buffer in bytes.

Type:

int

pc_structure

The structure of the push constant buffer.

Type:

List[vc.StructElement]

uniform_structure

The structure of the uniform buffer.

Type:

List[vc.StructElement]

binding_type_list

The list of binding types.

Type:

List[BindingType]

backend: CodeGenBackend
binding_access: List[Tuple[bool, bool]]
binding_type_list: List[BindingType]
body: str
exec_count_name: str | None
get_arg_names_and_defaults() → List[Tuple[str, Any]]
header: str
make_source(x: int, y: int, z: int) → str
name: str
pc_size: int
pc_structure: List[StructElement]
resource_binding_base: int
shader_arg_infos: List[ShaderArgumentInfo]
uniform_structure: List[StructElement]
class vkdispatch.codegen.ShaderFlags(*values)

Bases: IntFlag

NONE = 0
NO_EXEC_BOUNDS = 4
NO_PRINTF = 2
NO_SUBGROUP_OPS = 1
class vkdispatch.codegen.ShaderVariable(var_type: dtype, name: str | None = None, raw_name: str | None = None, lexical_unit: bool = False, settable: bool = False, register: bool = False, parents: List[ShaderVariable] = None, is_conjugate: bool = False, buffer_root: ShaderVariable | None = None, buffer_index_expr: str | None = None)

Bases: BaseVariable

buffer_index_expr: str | None
buffer_root: ShaderVariable | None
conjugate() → ShaderVariable
is_complex: bool
is_conjugate: bool | None
set_value(value: ShaderVariable) → None
swizzle(components: str) → ShaderVariable
to_dtype(var_type: dtype) → ShaderVariable
to_register(var_name: str = None) → ShaderVariable
class vkdispatch.codegen.StructElement(name: str, dtype: dtype, count: int)

Bases: object

A dataclass that represents an element of a struct.

name

The name of the element.

Type:

str

dtype

The dtype of the element.

Type:

vd.dtype

count

The count of the element.

Type:

int

count: int
dtype: dtype
name: str
vkdispatch.codegen.Var

alias of Variable

class vkdispatch.codegen.Variable

Bases: ShaderVariable, Generic[_ArgType]

vkdispatch.codegen.abs(var: Any) → ShaderVariable | float
vkdispatch.codegen.acos(var: Any) → ShaderVariable | float
vkdispatch.codegen.acosh(var: Any) → ShaderVariable | float
vkdispatch.codegen.all(*args: List[ShaderVariable])
vkdispatch.codegen.any(*args: List[ShaderVariable])
vkdispatch.codegen.asin(var: Any) → ShaderVariable | float
vkdispatch.codegen.asinh(var: Any) → ShaderVariable | float
vkdispatch.codegen.atan(var: Any) → ShaderVariable | float
vkdispatch.codegen.atan2(y: Any, x: Any) → ShaderVariable | float
vkdispatch.codegen.atanh(var: Any) → ShaderVariable | float
vkdispatch.codegen.atomic_add(mem: ShaderVariable, y: Any) → ShaderVariable
vkdispatch.codegen.barrier()
vkdispatch.codegen.c128

alias of _CF128

vkdispatch.codegen.c32

alias of _CF32

vkdispatch.codegen.c64

alias of _CF64

vkdispatch.codegen.ceil(var: Any) → ShaderVariable | float
vkdispatch.codegen.clamp(x: Any, min_val: Any, max_val: Any) → ShaderVariable | float
vkdispatch.codegen.clip(x: Any, min_val: Any, max_val: Any) → ShaderVariable | float
vkdispatch.codegen.comment(comment: str, preceding_new_line: bool = True) → None
vkdispatch.codegen.complex_from_euler_angle(angle: ShaderVariable)
vkdispatch.codegen.cos(var: Any) → ShaderVariable | float
vkdispatch.codegen.cosh(var: Any) → ShaderVariable | float
vkdispatch.codegen.cross(x: ShaderVariable, y: ShaderVariable) → ShaderVariable
vkdispatch.codegen.degrees(var: Any) → ShaderVariable | float
vkdispatch.codegen.determinant(var: ShaderVariable) → ShaderVariable
vkdispatch.codegen.distance(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.dot(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.dv2

alias of _V2F64

vkdispatch.codegen.dv3

alias of _V3F64

vkdispatch.codegen.dv4

alias of _V4F64

vkdispatch.codegen.else_block()
vkdispatch.codegen.else_if_block(arg: ShaderVariable)
vkdispatch.codegen.exp(var: Any) → ShaderVariable | float
vkdispatch.codegen.exp2(var: Any) → ShaderVariable | float
vkdispatch.codegen.f16

alias of _F16

vkdispatch.codegen.f32

alias of _F32

vkdispatch.codegen.f64

alias of _F64

vkdispatch.codegen.float_bits_to_int(var: Any) → ShaderVariable | int
vkdispatch.codegen.float_bits_to_uint(var: Any) → ShaderVariable | int
vkdispatch.codegen.floor(var: Any) → ShaderVariable | float
vkdispatch.codegen.fma(a: Any, b: Any, c: Any) → ShaderVariable | float
vkdispatch.codegen.fract(var: Any) → ShaderVariable | float
vkdispatch.codegen.get_builder() → ShaderBuilder
vkdispatch.codegen.get_codegen_backend() → CodeGenBackend
vkdispatch.codegen.get_shader_print_line_numbers() → bool
vkdispatch.codegen.global_invocation_id()
vkdispatch.codegen.group_memory_barrier()
vkdispatch.codegen.hv2

alias of _V2F16

vkdispatch.codegen.hv3

alias of _V3F16

vkdispatch.codegen.hv4

alias of _V4F16

vkdispatch.codegen.i16

alias of _I16

vkdispatch.codegen.i32

alias of _I32

vkdispatch.codegen.i64

alias of _I64

vkdispatch.codegen.if_block(arg: ShaderVariable)
vkdispatch.codegen.ihv2

alias of _V2I16

vkdispatch.codegen.ihv3

alias of _V3I16

vkdispatch.codegen.ihv4

alias of _V4I16

vkdispatch.codegen.inf_f16()
vkdispatch.codegen.inf_f32()
vkdispatch.codegen.inf_f64()
vkdispatch.codegen.int_bits_to_float(var: Any) → ShaderVariable | float
vkdispatch.codegen.inverse(var: ShaderVariable) → ShaderVariable
vkdispatch.codegen.inversesqrt(var: Any) → ShaderVariable | float
vkdispatch.codegen.isinf(var: Any) → ShaderVariable | bool
vkdispatch.codegen.isnan(var: Any) → ShaderVariable | bool
vkdispatch.codegen.iv2

alias of _V2I32

vkdispatch.codegen.iv3

alias of _V3I32

vkdispatch.codegen.iv4

alias of _V4I32

vkdispatch.codegen.length(var: Any) → ShaderVariable | float
vkdispatch.codegen.local_invocation_id()
vkdispatch.codegen.local_invocation_index()
vkdispatch.codegen.log(var: Any) → ShaderVariable | float
vkdispatch.codegen.log2(var: Any) → ShaderVariable | float
vkdispatch.codegen.m2

alias of _M2F32

vkdispatch.codegen.m3

alias of _M3F32

vkdispatch.codegen.m4

alias of _M4F32

vkdispatch.codegen.matrix_comp_mult(x: ShaderVariable, y: ShaderVariable) → ShaderVariable
vkdispatch.codegen.max(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.memory_barrier()
vkdispatch.codegen.memory_barrier_buffer()
vkdispatch.codegen.memory_barrier_image()
vkdispatch.codegen.memory_barrier_shared()
vkdispatch.codegen.min(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.mix(x: Any, y: Any, a: Any) → ShaderVariable | float
vkdispatch.codegen.mod(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.modf(x: Any, y: Any) → Tuple[ShaderVariable, ShaderVariable]
vkdispatch.codegen.mult_complex(arg1: ShaderVariable, arg2: ShaderVariable)
vkdispatch.codegen.new_complex128_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_complex32_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_complex64_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_complex_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_dvec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_dvec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_dvec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_float16_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_float64_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_float_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_hvec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_hvec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_hvec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ihvec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ihvec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ihvec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_int16_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_int64_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_int_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ivec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ivec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_ivec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_mat2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_mat3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_mat4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_register(var_type: dtype, *args, var_name: str | None = None)
vkdispatch.codegen.new_uhvec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uhvec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uhvec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uint16_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uint64_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uint_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uvec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uvec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_uvec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_vec2_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_vec3_register(*args, var_name: str | None = None)
vkdispatch.codegen.new_vec4_register(*args, var_name: str | None = None)
vkdispatch.codegen.ninf_f16()
vkdispatch.codegen.ninf_f32()
vkdispatch.codegen.ninf_f64()
vkdispatch.codegen.normalize(var: ShaderVariable) → ShaderVariable
vkdispatch.codegen.num_subgroups()
vkdispatch.codegen.num_workgroups()
vkdispatch.codegen.outer_product(x: ShaderVariable, y: ShaderVariable) → ShaderVariable
vkdispatch.codegen.pow(x: Any, y: Any) → ShaderVariable | float
vkdispatch.codegen.print(*args: Any, separator=' ')
vkdispatch.codegen.printf(format: str, *args: Any)
vkdispatch.codegen.radians(var: Any) → ShaderVariable | float
vkdispatch.codegen.ravel_index(index: ShaderVariable | int, shape: ShaderVariable | Tuple[int, ...])
vkdispatch.codegen.return_statement(arg=None)
vkdispatch.codegen.round(var: Any) → ShaderVariable | float
vkdispatch.codegen.round_even(var: Any) → ShaderVariable | float
vkdispatch.codegen.scope_block(indent: bool = True, comment: str = None)
vkdispatch.codegen.set_builder(builder: ShaderBuilder)
vkdispatch.codegen.set_codegen_backend(backend: CodeGenBackend | str | None)
vkdispatch.codegen.set_shader_print_line_numbers(value: bool)
vkdispatch.codegen.shader_context(flags: ~vkdispatch.codegen.builder.ShaderFlags = <ShaderFlags.NONE: 0>)
vkdispatch.codegen.shared_buffer(var_type: dtype, size: int, var_name: str | None = None)
vkdispatch.codegen.sign(var: Any) → ShaderVariable | float
vkdispatch.codegen.sin(var: Any) → ShaderVariable | float
vkdispatch.codegen.sinh(var: Any) → ShaderVariable | float
vkdispatch.codegen.smoothstep(edge0: Any, edge1: Any, x: Any) → ShaderVariable | float
vkdispatch.codegen.sqrt(var: Any) → ShaderVariable | float
vkdispatch.codegen.step(edge: Any, x: Any) → ShaderVariable | float
vkdispatch.codegen.str_to_dtype(var_type: dtype, value: str, parents: list | None = None, lexical_unit: bool = False, settable: bool = False, register: bool = False)
vkdispatch.codegen.subgroup_add(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_and(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_barrier()
vkdispatch.codegen.subgroup_elect()
vkdispatch.codegen.subgroup_id()
vkdispatch.codegen.subgroup_invocation_id()
vkdispatch.codegen.subgroup_max(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_min(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_mul(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_or(arg1: ShaderVariable)
vkdispatch.codegen.subgroup_size()
vkdispatch.codegen.subgroup_xor(arg1: ShaderVariable)
vkdispatch.codegen.tan(var: Any) → ShaderVariable | float
vkdispatch.codegen.tanh(var: Any) → ShaderVariable | float
vkdispatch.codegen.to_complex(*args)
vkdispatch.codegen.to_complex128(*args)
vkdispatch.codegen.to_complex32(*args)
vkdispatch.codegen.to_complex64(*args)
vkdispatch.codegen.to_dtype(var_type: dtype, *args)
vkdispatch.codegen.to_dvec2(*args)
vkdispatch.codegen.to_dvec3(*args)
vkdispatch.codegen.to_dvec4(*args)
vkdispatch.codegen.to_float(*args)
vkdispatch.codegen.to_float16(*args)
vkdispatch.codegen.to_float64(*args)
vkdispatch.codegen.to_hvec2(*args)
vkdispatch.codegen.to_hvec3(*args)
vkdispatch.codegen.to_hvec4(*args)
vkdispatch.codegen.to_ihvec2(*args)
vkdispatch.codegen.to_ihvec3(*args)
vkdispatch.codegen.to_ihvec4(*args)
vkdispatch.codegen.to_int(*args)
vkdispatch.codegen.to_int16(*args)
vkdispatch.codegen.to_int64(*args)
vkdispatch.codegen.to_ivec2(*args)
vkdispatch.codegen.to_ivec3(*args)
vkdispatch.codegen.to_ivec4(*args)
vkdispatch.codegen.to_mat2(*args)
vkdispatch.codegen.to_mat3(*args)
vkdispatch.codegen.to_mat4(*args)
vkdispatch.codegen.to_uhvec2(*args)
vkdispatch.codegen.to_uhvec3(*args)
vkdispatch.codegen.to_uhvec4(*args)
vkdispatch.codegen.to_uint(*args)
vkdispatch.codegen.to_uint16(*args)
vkdispatch.codegen.to_uint64(*args)
vkdispatch.codegen.to_uvec2(*args)
vkdispatch.codegen.to_uvec3(*args)
vkdispatch.codegen.to_uvec4(*args)
vkdispatch.codegen.to_vec2(*args)
vkdispatch.codegen.to_vec3(*args)
vkdispatch.codegen.to_vec4(*args)
vkdispatch.codegen.transpose(var: ShaderVariable) → ShaderVariable
vkdispatch.codegen.trunc(var: Any) → ShaderVariable | float
vkdispatch.codegen.u16

alias of _U16

vkdispatch.codegen.u32

alias of _U32

vkdispatch.codegen.u64

alias of _U64

vkdispatch.codegen.uhv2

alias of _V2U16

vkdispatch.codegen.uhv3

alias of _V3U16

vkdispatch.codegen.uhv4

alias of _V4U16

vkdispatch.codegen.uint_bits_to_float(var: Any) → ShaderVariable | float
vkdispatch.codegen.unravel_index(index: ShaderVariable | Tuple[int, ...], shape: ShaderVariable | Tuple[int, ...])
vkdispatch.codegen.uv2

alias of _V2U32

vkdispatch.codegen.uv3

alias of _V3U32

vkdispatch.codegen.uv4

alias of _V4U32

vkdispatch.codegen.v2

alias of _V2F32

vkdispatch.codegen.v3

alias of _V3F32

vkdispatch.codegen.v4

alias of _V4F32

vkdispatch.codegen.while_block(arg: ShaderVariable)
vkdispatch.codegen.workgroup_id()
vkdispatch.codegen.workgroup_size()