Module Sarek_vulkan.Vulkan_api_kernel

module Device = Vulkan_api_device
module Memory = Vulkan_api_memory
module Stream = Vulkan_api_stream
type t = {
  1. shader_module : Vulkan_types.vk_shader_module;
  2. pipeline : Vulkan_types.vk_pipeline;
  3. pipeline_layout : Vulkan_types.vk_pipeline_layout;
  4. descriptor_set_layout : Vulkan_types.vk_descriptor_set_layout;
  5. descriptor_pool : Vulkan_types.vk_descriptor_pool;
  6. descriptor_set : Vulkan_types.vk_descriptor_set;
  7. name : string;
  8. num_bindings : int;
  9. device : Device.t;
}
type arg =
  1. | ArgBuffer : _ Memory.buffer -> arg
  2. | ArgInt32 : int32 -> arg
  3. | ArgInt64 : int64 -> arg
  4. | ArgFloat32 : float -> arg
  5. | ArgFloat64 : float -> arg
  6. | ArgPtr : nativeint -> arg
type any_buffer =
  1. | AnyBuf : 'a Memory.buffer -> any_buffer

Existential wrapper to hide buffer type parameter

type scalar_arg =
  1. | SInt32 of int32
  2. | SInt64 of int64
  3. | SFloat32 of float
  4. | SFloat64 of float

Scalar kernel argument, tagged with its byte width/value so the push-constant block can be materialized in one pass at launch time (see build_push_constants) instead of being appended byte-by-byte in whatever order set_arg_* happens to be called.

type args = {
  1. buffer_store : any_buffer Spoc_framework.Kernel_args.t;
  2. scalar_store : scalar_arg Spoc_framework.Kernel_args.t;
  3. mutable descriptor_set : Vulkan_types.vk_descriptor_set;
}
val resolve_bindings : args -> (int * any_buffer) list

Vulkan descriptor binding numbers are assigned to buffer-typed kernel parameters only, in the order those parameters appear in the kernel signature (see the GLSL codegen, which numbers `layout(binding = N)` this way, skipping scalar parameters entirely). Since kernel-arg idx is shared across buffers and scalars, idx itself is not usable as a binding number directly; instead the buffer with the Nth-smallest idx among buffer args gets binding N. Because buffers cannot be reordered by the caller relative to each other without also changing which kernel parameter they represent, this rank-by-idx is stable and order-of-call-independent, unlike the previous per-call sequential counter it replaces.

val validate_buffer_indices : expected_count:int -> 'a Spoc_framework.Kernel_args.t -> (unit, string) Stdlib.result

Validate the set of buffer-argument indices actually bound in store against the number of buffer bindings the compiled kernel expects (expected_count, derived from the GLSL source's `binding = N` declarations \- see num_bindings in compile).

resolve_bindings above silently compresses whatever indices are present into dense, ascending descriptor bindings; that rank-mapping is correct for a valid index *set*, but gives no signal at all if the caller passed a nonsensical one (negative indices, or simply the wrong number of buffers - e.g. a caller-side idx typo that drops one buffer and duplicates another slot). This checks the set's sanity before resolve_bindings's compression can paper over it. Exposed for direct unit testing (no hardware required).

val cache : (string, t) Spoc_framework.Guarded_cache.t
val create_shader_module : Device.t -> string -> Vulkan_types.vk_shader_module

Create shader module from SPIR-V

val compile : Device.t -> name:string -> source:string -> t

Compile GLSL source to compute pipeline

val compile_cached : Device.t -> name:string -> source:string -> t
val clear_cache : unit -> unit
val create_args : unit -> args
val set_arg_buffer : args -> int -> 'a Memory.buffer -> unit
val set_arg_int32 : args -> int -> int32 -> unit
val set_arg_int64 : args -> int -> int64 -> unit
val set_arg_float32 : args -> int -> float -> unit
val set_arg_float64 : args -> int -> float -> unit
val set_arg_ptr : 'a -> 'b -> 'c -> 'd
val push_constant_limit : int
val build_push_constants : args -> bytes option

Materialize the push-constant byte block for args at launch time, from the full ordered argument set (buffers + scalars), matching EXACTLY the GLSL block layout emitted by Sarek_ir_glsl.gen_push_constants (sarek/codegen/Sarek_ir_glsl.ml:889-919): all vector lengths first, in vector-declaration order, followed by all user scalar parameters, in declaration order.

  • "Vector-declaration order" here is the same order used for descriptor bindings: the buffer with the Nth-smallest caller-supplied idx gets the Nth length slot (see resolve_bindings above / gen_push_constants's vectors list, which is built by iterating kernel params in order).
  • "Scalar declaration order" is the ascending order of the scalar arguments' own caller-supplied idx (see gen_push_constants's scalars list, same iteration).

This must not be done incrementally inside each set_arg_* call: the caller may invoke those in any call order (e.g. a vector, then a scalar, then another vector), which does not match the GLSL grouping, so the block can only be assembled correctly once every argument is known.

val launch : t -> args:args -> grid:Spoc_framework.Framework_sig.dims -> block:Spoc_framework.Framework_sig.dims -> shared_mem:'a -> stream:Stream.t option -> unit