Module Sarek_native.Native_plugin_base

val native_kernels : (string, Spoc_framework.Framework_sig.exec_arg array -> (int * int * int) -> (int * int * int) -> unit) Stdlib.Hashtbl.t

Registry for native kernel functions.

Kernels are registered by name. The function signature uses exec_arg array for type-safe argument passing:

  • args: Framework_sig.exec_arg array (typed kernel arguments)
  • grid: (gx, gy, gz) grid dimensions
  • block: (bx, by, bz) block dimensions

Defined at module level so it can be accessed by both the Native module and the external registration functions.

module Native : sig ... end
val init : unit -> unit
val register_kernel : string -> (Spoc_framework.Framework_sig.exec_arg array -> (int * int * int) -> (int * int * int) -> unit) -> unit

Register a native kernel function by name.

This is called by PPX-generated code at module load time to register the pre-compiled OCaml kernel function. The function signature matches what Sarek_cpu_runtime expects:

  • args: Framework_sig.exec_arg array of kernel arguments
  • grid: (gx, gy, gz) grid dimensions
  • block: (bx, by, bz) block dimensions

Example:

  let () =
    Sarek_native.Native_plugin.register_kernel
      "my_kernel"
      (fun args grid block ->
        Sarek_cpu_runtime.run_threadpool
          ~has_barriers:false
          ~block
          ~grid
          my_kernel_fn
          args)
val kernel_registered : string -> bool

Check if a kernel is registered

val list_kernels : unit -> string list

List all registered kernels (for debugging)

val run_kernel_raw : name:string -> args:Spoc_framework.Framework_sig.exec_arg array -> grid:(int * int * int) -> block:(int * int * int) -> unit

Run a kernel directly with typed arguments. This bypasses the Runtime API.

  • parameter name

    The registered kernel name

  • parameter args

    Array of exec_arg (typed kernel arguments)

  • parameter grid

    (gx, gy, gz) grid dimensions

  • parameter block

    (bx, by, bz) block dimensions