Module Sarek_env

module StringMap : sig ... end
type var_info = {
  1. vi_type : Sarek_types.typ;
  2. vi_mutable : bool;
  3. vi_is_param : bool;
    (*

    Kernel parameter?

    *)
  4. vi_index : int;
    (*

    For parameter ordering

    *)
  5. vi_is_vec : bool;
    (*

    Is this a vector parameter?

    *)
}

Information about a variable

type intrinsic_ref =
  1. | IntrinsicRef of string list * string
    (*

    module_path, function_name

    *)
  2. | CorePrimitiveRef of string
    (*

    Core primitive name

    *)

Reference to an intrinsic function or constant.

  • IntrinsicRef: Reference to a stdlib module intrinsic. The module path allows the PPX to generate correct references for compile-time type checking. For example:
  • IntrinsicRef ("Sarek_stdlib"; "Float32", "sin") -> Sarek_stdlib.Float32.sin
  • IntrinsicRef ("Sarek_stdlib"; "Int32", "add_int32") -> Sarek_stdlib.Int32.add_int32
  • CorePrimitiveRef: Reference to a core GPU primitive defined in Sarek_core_primitives. These have compile-time semantic properties (variance, convergence, purity) that cannot be overridden by user code. Device implementations are still resolved via Sarek_registry at JIT time.

This enables extensibility: user libraries can define their own intrinsics via %sarek_intrinsic and the PPX will reference them correctly.

val intrinsic_ref_name : intrinsic_ref -> string

Get the qualified name of an intrinsic ref for debugging/printing

val intrinsic_ref_display_name : intrinsic_ref -> string

Get user-friendly display name for error messages (no core prefix)

type intrinsic_fun_info = {
  1. intr_type : Sarek_types.typ;
  2. intr_ref : intrinsic_ref;
    (*

    Reference to stdlib module function

    *)
  3. intr_convergence : Sarek_core_primitives.convergence option;
    (*

    Convergence requirement. None means NoEffect. User libraries can specify ConvergencePoint or WarpConvergence.

    *)
}

Information about an intrinsic function. Device code is resolved at JIT time via Sarek_registry.

type intrinsic_const_info = {
  1. const_type : Sarek_types.typ;
  2. const_ref : intrinsic_ref;
    (*

    Reference to intrinsic constant

    *)
}

Information about an intrinsic constant. Device code is resolved at JIT time via Sarek_registry.

type type_info =
  1. | TIRecord of {
    1. ti_name : string;
    2. ti_fields : (string * Sarek_types.typ * bool) list;
      (*

      name, type, mutable

      *)
    }
  2. | TIVariant of {
    1. ti_name : string;
    2. ti_constrs : (string * Sarek_types.typ option) list;
      (*

      constructor name, optional arg type

      *)
    }

Information about a custom type

type t = {
  1. vars : var_info StringMap.t;
  2. types : type_info StringMap.t;
  3. intrinsic_funs : intrinsic_fun_info StringMap.t;
  4. intrinsic_consts : intrinsic_const_info StringMap.t;
  5. constructors : (string * Sarek_types.typ option) StringMap.t;
    (*

    constr -> (type_name, arg_type)

    *)
  6. fields : (string * int * Sarek_types.typ * bool) StringMap.t;
    (*

    field -> (type_name, index, type, mutable)

    *)
  7. current_level : int;
    (*

    For let-polymorphism

    *)
  8. local_funs : (string * Sarek_scheme.scheme) StringMap.t;
    (*

    Local function definitions with polymorphic schemes

    *)
}

The typing environment - immutable

val empty : t

Empty environment

Environment operations - all return new environments

val add_var : StringMap.key -> var_info -> t -> t
val add_type : StringMap.key -> type_info -> t -> t
val add_intrinsic_fun : StringMap.key -> intrinsic_fun_info -> t -> t
val add_intrinsic_const : StringMap.key -> intrinsic_const_info -> t -> t
val add_local_fun : StringMap.key -> Sarek_scheme.scheme -> t -> t

A kernel-local or [@sarek.module] function. If name is also a refused intrinsic (Sarek_unlowered.table), the user's function shadows it: name resolution consults intrinsics before local functions, so without this a helper called warp_size would resolve to the refused intrinsic and be refused, or worse, silently call it. Only refused names are shadowed here; other names keep their existing resolution order.

val find_var : StringMap.key -> t -> var_info option
val find_type : StringMap.key -> t -> type_info option
val find_intrinsic_fun : StringMap.key -> t -> intrinsic_fun_info option
val find_intrinsic_const : StringMap.key -> t -> intrinsic_const_info option
val find_constructor : StringMap.key -> t -> (string * Sarek_types.typ option) option
val find_field : StringMap.key -> t -> (string * int * Sarek_types.typ * bool) option
val find_local_fun : StringMap.key -> t -> (string * Sarek_scheme.scheme) option
val enter_level : t -> t

Scope management

val exit_level : t -> t
val short_module_name : string list -> string

Get the short module name from a module path. E.g., "Sarek_stdlib"; "Float32" -> "Float32"

val open_module : string list -> t -> t

Open a module: bring its bindings into scope under short names. E.g., open_module "Float32" brings Float32.sin -> sin

Also handles legacy aliases:

  • Std -> Gpu (backward compatibility)
val with_stdlib : t -> t

Create standard library environment from core primitives and the PPX registry.

Step 1: Add all core primitives from Sarek_core_primitives. These have compile-time semantic properties (variance, convergence, purity) that the PPX uses for analysis. Core primitives use CorePrimitiveRef.

Step 2: Add library intrinsics from Sarek_ppx_registry. These may shadow core primitives if they have the same name (which provides device implementations). Library intrinsics use IntrinsicRef.

Intrinsics are registered under module-qualified names (e.g., "Float32.sin") to avoid ambiguity between Float32.sqrt and Float64.sqrt. Use `open_module` to bring a module's bindings into scope under short names.

IMPORTANT: The caller must ensure stdlib modules are initialized before calling this function (e.g., via Sarek_stdlib.force_init()). This is done in Sarek_ppx.ml to avoid circular dependencies.

type lookup_result =
  1. | LVar of var_info
  2. | LIntrinsicConst of intrinsic_const_info
  3. | LIntrinsicFun of intrinsic_fun_info
  4. | LConstructor of string * Sarek_types.typ option
    (*

    type_name, arg_type

    *)
  5. | LLocalFun of string * Sarek_scheme.scheme
    (*

    name and type scheme (for polymorphism)

    *)
  6. | LNotFound

Lookup that checks all namespaces for an identifier

val lookup : StringMap.key -> t -> lookup_result
val known_names_in_module_of : string -> t -> Stdlib.String.t list

For a qualified name M.f, the qualified names the environment binds under the same prefix M., sorted (used to suggest a fix for a typo or a module whose intrinsics are only partly linked). [] for an unqualified name.

val pp_env : Stdlib.Format.formatter -> t -> unit

Debug: print environment contents