Module Sarek_codegen.Sarek_ir_wgsl

module Codegen_error : sig ... end

Local error module — tagged as "WebGPU" in error messages.

val current_variants : (string * (string * Sarek_ir_types.elttype list) list) list Stdlib.ref

Current kernel's variant definitions (set during generate)

val current_framework : string option Stdlib.ref

Current framework name — mirrors the other generators so set_framework in Sarek_transpile can set it, and the pure registry resolves Float32 math with framework="WGSL" (falls through to the generic sin spelling).

Type Mapping

val mangle_name : string -> string
val wgsl_reserved_keywords : string list

WGSL reserved keywords that cannot be used as identifiers

val escape_wgsl_name : string -> string

Escape reserved WGSL keywords by adding 'v' suffix.

val wgsl_type_of_elttype : Sarek_ir_types.elttype -> string

Map Sarek IR element type to WGSL type string. Float64 (f64) is not supported in WebGPU — callers must check for TFloat64 before reaching this function and raise Codegen_error.unsupported_construct.

val has_float64 : Sarek_ir_types.elttype -> bool

Check whether an elttype (recursively) uses Float64.

val wgsl_thread_intrinsic : string -> string

Thread Intrinsics

WGSL uses three distinct builtins:

  • local_invocation_id (sarek_lid) — thread within workgroup
  • workgroup_id (sarek_wid) — workgroup index in the dispatch grid
  • global_invocation_id (sarek_gid) — globally unique thread index

All are vec3<u32>; we cast to i32 to match the IR's i32 type for thread ids. The entry point declares all three builtins; unused ones are harmless (WGSL permits unused builtin params).

Expression Generation

val scalar_param_names : string list Stdlib.ref

Names of scalar kernel params — accessed as params.<name> in WGSL.

val gen_expr : Stdlib.Buffer.t -> Sarek_ir_types.expr -> unit
val gen_binop : Sarek_ir_types.binop -> string
val gen_unop : Sarek_ir_types.unop -> string
val gen_intrinsic : Stdlib.Buffer.t -> string list -> string -> Sarek_ir_types.expr list -> unit

L-value Generation

val gen_lvalue : Stdlib.Buffer.t -> Sarek_ir_types.lvalue -> unit

Statement Generation

val indent_nested : string -> string
val gen_match_pattern : Stdlib.Buffer.t -> string -> string -> string -> string list -> (string -> Sarek_ir_types.elttype list option) -> unit
val gen_var_decl : Stdlib.Buffer.t -> string -> mutable_:bool -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.expr -> unit
val gen_stmt : Stdlib.Buffer.t -> string -> Sarek_ir_types.stmt -> unit

Helper Function Generation

val gen_helper_func : Stdlib.Buffer.t -> Sarek_ir_types.helper_func -> unit

Record and Variant Type Generation

val gen_record_def : Stdlib.Buffer.t -> (string * (string * Sarek_ir_types.elttype) list) -> unit
val gen_variant_def : Stdlib.Buffer.t -> (string * (string * Sarek_ir_types.elttype list) list) -> unit

Emit a WGSL variant type. WGSL has no enums or unions. We emit:

  • const <CNAME> : i32 = N; for each constructor tag
  • a struct with tag : i32 and flat payload fields
  • fn make_<Type>_<Constr>(...) -> <Type> constructors

Buffer / Uniform Binding Generation

val collect_workgroup_decls : Sarek_ir_types.stmt -> (string * Sarek_ir_types.elttype * Sarek_ir_types.expr) list

Collect workgroup shared array declarations from a statement tree.

val gen_workgroup_module_decls : Stdlib.Buffer.t -> (string * Sarek_ir_types.elttype * Sarek_ir_types.expr) list -> unit
val split_params : Sarek_ir_types.decl list -> Sarek_ir_types.var list * Sarek_ir_types.var list

Separate kernel params into vectors (storage buffers) and scalars (uniform).

val gen_bindings : Stdlib.Buffer.t -> Sarek_ir_types.decl list -> string list

Emit storage buffer bindings and the Params uniform struct. Returns the list of scalar param names (for scalar_param_names ref).

Main generate functions

val wgsl_header : kernel_name:string -> ?block:(int * int * int) -> unit -> string
val params_have_float64 : Sarek_ir_types.decl list -> bool

Check if any kernel param uses Float64.

val generate : ?block:(int * int * int) -> ?log:(string -> unit) -> Sarek_ir_types.kernel -> string

Generate complete WGSL source for a kernel.

val generate_with_types : ?block:(int * int * int) -> ?log:(string -> unit) -> types:(string * (string * Sarek_ir_types.elttype) list) list -> Sarek_ir_types.kernel -> string

Generate WGSL source with custom type definitions.

ABI descriptor

val abi : ?block:(int * int * int) -> Sarek_ir_types.kernel -> Sarek_wgsl_abi.t

Build the ABI descriptor for a kernel. Reuses split_params and escape_wgsl_name / wgsl_type_of_elttype so the descriptor cannot drift from gen_bindings.

Raises Codegen_error.unsupported_construct for f64 parameters (same error as generate).