Module Sarek_registry

type type_info = {
  1. ti_name : string;
  2. ti_device : string -> string;
  3. ti_size : int;
}

Information about a primitive/intrinsic type (float32, int64, etc.)

type field_info = {
  1. field_name : string;
  2. field_type : string;
  3. field_mutable : bool;
}

Information about a record field

type record_info = {
  1. ri_name : string;
  2. ri_fields : field_info list;
  3. ri_size : int;
}

Information about a record type (user-defined via @@sarek.type)

type constructor_info = {
  1. ctor_name : string;
  2. ctor_arg_type : string option;
}

Information about a variant constructor

type variant_info = {
  1. vi_name : string;
  2. vi_constructors : constructor_info list;
}

Information about a variant type

type fun_info = {
  1. fi_name : string;
  2. fi_arity : int;
  3. fi_device : string -> string;
  4. fi_arg_types : string list;
  5. fi_ret_type : string;
}

Information about an intrinsic function

val type_registry : (string, type_info) Stdlib.Hashtbl.t

Type registry - maps type names to their info (primitives)

val record_registry : (string, record_info) Stdlib.Hashtbl.t

Record registry - maps type names to their info (user-defined records)

val variant_registry : (string, variant_info) Stdlib.Hashtbl.t

Variant registry - maps type names to their info (user-defined variants)

val fun_registry : (string list * string, fun_info) Stdlib.Hashtbl.t

Function registry - maps (module_path, name) to their info

val register_type : string -> device:(string -> string) -> size:int -> unit

Register a primitive type

val register_record : string -> fields:field_info list -> size:int -> unit

Register a record type (called by PPX-generated code for @@sarek.type)

val register_variant : string -> constructors:constructor_info list -> unit

Register a variant type (called by PPX-generated code for @@sarek.type)

val register_fun : ?module_path:string list -> string -> arity:int -> device:(string -> string) -> arg_types:string list -> ret_type:string -> unit

Register an intrinsic function

type host_value =
  1. | HInt32 of int32
  2. | HInt64 of int64
  3. | HFloat32 of float
  4. | HFloat64 of float
  5. | HBool of bool
  6. | HUnit

A scalar crossing between an evaluator and an intrinsic's host (OCaml) implementation. One constructor per scalar type an intrinsic signature can name: int32, int64, float32 (an OCaml float holding a binary32 value), float/float64, bool and unit.

val host_registry : (string list * string, host_value list -> host_value) Stdlib.Hashtbl.t

Host implementations of intrinsic functions, keyed like fun_registry. let%sarek_intrinsic registers one for every function whose signature is all scalars; it is what lets an evaluator that has no built-in table entry for an intrinsic (the interpreter, for an intrinsic declared outside sarek's stdlib) run it. The implementation receives its arguments in the declared order, each with the constructor of its declared type, and raises Invalid_argument on any other shape.

val register_host_impl : ?module_path:string list -> string -> (host_value list -> host_value) -> unit
val find_host_impl : ?module_path:string list -> string -> (host_value list -> host_value) option
val find_type : string -> type_info option

Find a primitive type by name

val find_record : string -> record_info option

Find a record type by name

val find_variant : string -> variant_info option

Find a variant type by name

val find_fun : ?module_path:string list -> string -> fun_info option

Find a function by name, optionally in a module

val is_type : string -> bool

Check if a name is a registered primitive type

val is_record : string -> bool

Check if a name is a registered record type

val is_variant : string -> bool

Check if a name is a registered variant type

val is_fun : ?module_path:string list -> string -> bool

Check if a name is a registered function

val type_device_code : string -> string -> string

Get device code for a type

val fun_device_code : ?module_path:string list -> string -> string -> string

Get device code for a function

val fun_device_template : ?module_path:string list -> ?framework:string -> string -> string option

Get the device-code template for a function on a given backend. This is for V2 IR codegens that don't have SPOC device objects.

framework is the caller's backend tag ("CUDA", "OpenCL", "Metal", …). It used to be hardcoded to "generic", which cuda_or_opencl resolves to the CUDA branch — so every backend reaching this fallback got the CUDA spelling, and an OpenCL or Metal kernel calling e.g. Float32.abs_float was emitted as fabsf(...). Neither OpenCL C nor MSL declares fabsf; both spell it fabs. The stdlib already declares both spellings (dev "fabsf" "fabs"); only this lookup was discarding the caller's framework.

?framework defaults to "generic" so existing non-backend callers keep the previous behaviour.

val find_record_by_short_name : string -> record_info option

Find a record by short name (last component after '.'). This handles cases where the registry uses qualified names like "Module.typename" but the custom_type uses just "typename".

val record_fields : string -> field_info list

Get record field info - tries exact match first, then short name

val variant_constructors : string -> constructor_info list

Get variant constructors

exception No_device_spelling of {
  1. framework : string;
  2. cuda : string;
}

Raised when a shading-language framework asks this two-way dispatch for a spelling it cannot supply. See cuda_or_opencl.

val cuda_or_opencl : string -> string -> string -> string