Module Sarek_codegen.Sarek_ir_wgsl

module Codegen_error : sig ... end

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

module Dispatch = Sarek_ir_intrinsic_dispatch
val bad_arity : string -> int -> int -> 'a

Raise a located invalid-argument-count error (atomic-arity helper for the shared Dispatch.emit_atomic).

type state = {
  1. variants : (string * (string * Sarek_ir_types.elttype list) list) list;
  2. scalar_params : string list;
}

Everything one run of generate_with_types needs to know that is not reachable from the IR node it is currently emitting. It is a VALUE threaded through the emit functions, not module state, and that is the whole point of backlog-185/200: three module-level refs used to hold this kind of state, so a second generation — on another domain, or simply a later one after Sarek_transpile had written current_framework — read the first one's values. Only two of those three refs have a successor field below; the third, current_framework, was retired with no replacement field at all.

The framework tag is NOT among them. It is the constant "WGSL", read at the one Dispatch.framework thunk in wgsl_backend and nowhere else in this emitter — in particular it does not reach Sarek_registry, because the post_hook that would wire that registry is deliberately inert here (see the comment on it). An earlier draft of backlog-185/200 threaded it as a ?framework argument; that was dropped, because no caller ever passed anything but "WGSL" and an untested parameter on a public API looks like protection without being any.

variants is the kernel's own kern_variants, read by the SMatch arm to recover a constructor's payload types. Derived from the kernel, so it could be re-derived at each use site; it is carried here because that is where the ref it replaces was read from.

scalar_params is the raw (unescaped) names of the kernel's scalar params, which live in the Params uniform and are therefore read as params.<name> rather than as a bare identifier — the one thing gen_expr's EVar arm cannot decide from the node alone. It is the only field that varies within one generation: gen_helper_func emits each helper under a smaller set, because the helper's own formals shadow like-named kernel scalars lexically. That narrowing is now a rebind of this field for the nested emit, so the enclosing generation's value is simply never touched.

val state : scalar_params:string list -> Sarek_ir_types.kernel -> state

The state for emitting k. scalar_params is a parameter rather than derived from k because the authority on which params became Params fields is gen_bindings, which emits them; the sole caller passes its result, so the emitted struct and the set gen_expr tests against are the same list.

Type Mapping

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

Names that must never be emitted as a WGSL identifier.

Three groups, kept separate because they are reserved for different reasons and a reader needs to know which ones are load-bearing:

  • Predeclared names (types, address spaces, access modes, attribute names). WGSL does accept these as user identifiers — they live in a scope outside module scope and are shadowable — so escaping them is defensive rather than mandatory: a kernel variable named f32 parses, but then f32 can no longer be spelled as a type in the same scope and the emitter would produce a shader that fails much later and much less legibly.
  • Keywords (WGSL §2.3 "Keyword Summary"). Hard parse errors.
  • Reserved words (WGSL §2.4 "Reserved Words"). Also hard parse errors, and this is where the previous table was materially incomplete: plausible OCaml variable names such as ref, set, get, from, shared, filter, target and where are all reserved in WGSL and were being emitted verbatim, producing a shader that no WebGPU implementation accepts.

The keyword and reserved-word groups are not a transcription from memory: every entry below was verified to be actually rejected by running naga 30.0.0 (the validator ci/assert-toolchain.sh pins) over a minimal compute shader declaring var <name> : i32. The probe also established that the predeclared group is not rejected, which is why it is labelled defensive above rather than merged into the other two. Re-run the probe when bumping naga: a word moving between these groups changes nothing (both are escaped), but a word being added to WGSL must be added here.

val wgsl_escape_prefix : string

Escape identifiers that WGSL forbids.

Four things make a name unusable as a WGSL identifier. All four are handled by the same rewrite — prefixing with "sarek_" — for the injectivity reason set out below.

  • Reserved prefix (WGSL §4.4.1): an identifier must not start with a double underscore. The frontend's tail-recursion elimination (Sarek_tailrec_elim) renames every eliminated loop parameter to "__" ^ name, and the native backend emits "__v%d"/"__m%d" temporaries, so any kernel whose recursion is turned into a loop reached this emitter with a __-prefixed variable. C-family targets (CUDA, OpenCL, GLSL, Metal) accept those names, WGSL rejects them outright at parse time ("Identifier starts with a reserved prefix"), which made the emitted shader unusable on every WebGPU implementation. Prefixing keeps the name recognisable and moves the double underscore off the front, where it is legal.
  • Bare underscore (same clause): a lone "_" is not an identifier in WGSL either — it is the phony-assignment target. An OCaml wildcard or generated placeholder reaching the emitter as "_" would produce let _ : i32 = ..., which naga rejects.
  • Reserved names: anything in wgsl_reserved_keywords — WGSL keywords and reserved words, plus the predeclared and internal names this emitter escapes defensively.
  • The escaped namespace itself: a source name that already starts with "sarek_" is escaped too, so it cannot be confused with the image of one that was. This is what makes the rewrite injective; see below.

Injectivity

An earlier form of this function rewrote each problem separately — "_" to "sarek_", a "__" prefix to "sarek" ^ name, a keyword to name ^ "v" — and was not injective: "__i" and "sarek__i" both emitted "sarek__i", "_" and "sarek_" both emitted "sarek_", and "if" and "ifv" both emitted "ifv". Two source variables colliding on one WGSL name is not a cosmetic problem: the second var declaration shadows the first in the same block, every later read silently resolves to the wrong binding, and the shader still compiles. A wrong answer with no diagnostic is the worst failure mode available here.

The rule below is a single unconditional one: the whole "sarek_" prefix is reserved, and any name that is reserved, or that could be confused with an escaped name, is prefixed with it.

That this is injective on source identifiers is checkable by cases rather than by inspection:

  • the escaping branch only ever emits names starting with "sarek_";
  • the identity branch only ever emits names not starting with "sarek_", because such a name would have taken the escaping branch;
  • the two images are therefore disjoint, and each branch is injective on its own (name ↦ "sarek_" ^ name and the identity).

Generator-produced names

wgsl_generated_prefixes is exempt, and the exemption is structural rather than a convenience. rename_scalar_shadowing_locals mints sarek_scalar_shadow_* names and puts them into the IR as ordinary variables, so they reach this function again on the way out. A name in the escaped namespace cannot be a fixed point of the rule above — that is what reserving the prefix means — so re-escaping produced sarek_sarek_scalar_shadow_width_1, and the alternative of choosing an internal name that is a fixed point is self-defeating: every fixed point is, by definition, reachable from the identical source identifier. Internal names must therefore be minted in final form and left alone.

The cost is one contrived residual: a user identifier spelled exactly like a generator-internal name (a local literally named sarek_scalar_shadow_width_1) is no longer pushed out of the namespace and can collide with the generated one. That is strictly smaller than what this function replaced, which collided on __i/sarek__i, _/sarek_ and if/ifv — three families of ordinary, plausible names.

The output is also always a legal WGSL identifier: it never starts with "__" (an escaped name starts with "sarek", and an unescaped one cannot start with "__" or it would have been escaped), it is never a bare "_", and it is never reserved (no WGSL keyword or reserved word starts with "sarek_").

Residual

abi builds the uniform-struct length fields as "sarek_" ^ escape_wgsl_name v ^ "_length", a second construction in the same namespace that this function cannot police. A vector named "if" and a scalar named "sarek_if_length" in one kernel still collide there. Closing it needs a length-prefixed encoding for that namespace, which changes the emitted ABI field names, so it is deliberately left for its own change rather than folded into this one.

val wgsl_generated_prefixes : string list

Prefixes of names this generator mints itself, in already-final form. They round-trip unchanged through escape_wgsl_name; see its "Generator-produced names" section for why the exemption is unavoidable rather than a shortcut. Anything added here must be a name a generator constructs, never a name that can arrive from source.

val escape_wgsl_name : string -> string
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 is_array_shaped_operand : Sarek_ir_types.expr -> bool

Whole-value equality on a vector/array-typed operand — TVec/TArray in the IR — is NOT refused by the frontend (backlog-217; see Sarek_types.is_uncomparable_operand_typ, which deliberately excludes them because src = dst on a pointer-shaped C-family value emits (src == dst), and clang and glslang both accept that). WGSL is the exception: measured, naga rejects the identical shape with "Incompatible operands: Equal(Array …, _)" — WGSL has no equality operator on the `array<T>`/`array<T, N>` type naga assigns both a vector kernel parameter and a local array to, unlike a C-family pointer. This predicate is deliberately narrow: it only matches an operand that reaches gen_expr as the WHOLE vector/array value (a bare EVar of TVec/TArray type). Indexing (`a.(i)`) lowers to EArrayRead/EArrayReadExpr, which yields a scalar element and is unaffected — comparing two elements is ordinary scalar equality and WGSL accepts it.

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

L-value Generation

val gen_lvalue : state -> 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 : state -> Stdlib.Buffer.t -> string -> mutable_:bool -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.expr -> unit
val gen_stmt : state -> Stdlib.Buffer.t -> string -> Sarek_ir_types.stmt -> unit
val rename_scalar_shadowing_locals : scalar_names:string list -> Sarek_ir_types.stmt -> Sarek_ir_types.stmt

Alpha-rename kernel-body binders whose name collides with a scalar kernel param.

Scalar params are accessed in the body as params.<name>; gen_expr decides this per-EVar by checking the emit state's scalar_params set, which is per-generation but flat — it carries no scope, so it cannot distinguish a reference to the param from a reference to a local of the same name. A local let width = … (or let mut width = …) that shadows a scalar param width therefore has every body reference to width wrongly emitted as params.width — reading the uniform instead of the local. For an immutable self-binding local (let width = params.width) this is accidentally correct; for a mutated shadowing local it is a silent wrong result (valid WGSL, no error): the declaration uses the bare name (var width : i32 = params.width;) so writes hit the local, but every read is redirected to the immutable uniform.

This mirrors the GLSL backend's Sarek_ir_glsl.rename_pc_shadowing_locals; both delegate the shared traversal to Sarek_ir_codegen.rename_shadowing_locals. Each colliding binder (and its in-scope references) is rewritten to a fresh sarek_scalar_shadow_* name that is not a scalar param, so gen_expr's scalar_params check never matches it. The initializer is evaluated in the outer scope, so it still expands to params.<name>, preserving semantics. Unlike GLSL there is no vector-length collision: both spell the length sarek_<arr>_length (EArrayLen), but WGSL emits it as the field access params.sarek_<arr>_length, hardcoded with a params. prefix independent of any local, so a local cannot alias it — the collision set is scalar params only. WGSL-only.

Main generate functions

Helper Function Generation

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

Emit the sarek_fmod C-fmod helper (f32; WGSL has no f64) when the kernel uses fmod. Replaces the earlier bare % lowering, which shared C-fmod's two divergences (both raised in review): the single-pass x - y*trunc(x/y) loses quotient precision for large |x/y|, and an infinite divisor yields NaN where C defines fmod(x, ±inf) = x.

The body is a bounded exact reduction by power-of-two scaling (identical in shape to the GLSL sarek_fmod helper): scale d = |y| up by ×2 to the largest |y|·2^k ≤ |x|, then walk back down subtracting whenever r ≥ d. Every ×2/×0.5 is exact and each subtraction runs with d ≤ r < 2d (exact by Sterbenz), so r is the bit-exact remainder magnitude; the loops are bounded by the f32 exponent span (~277 iterations). The dividend's sign is restored by a bit-level copy.

WGSL has no isnan/isinf; infinity is detected by a magnitude test against the largest finite f32 (0x1.fffffep+127). |y| = inf returns x (C-conformant). The genuine NaN-domain cases (y = 0, |x| = inf) are NOT expressible in WGSL's float model — WGSL cannot synthesise a NaN — so they return x purely to keep the reduction loop terminating; this is the one residual divergence from C, unavoidable in WGSL and documented as such.

The helper name is a fixed sarek_fmod, emitted verbatim at both its definition here and its call site in gen_expr. A user helper of the same name no longer clashes with it: escape_wgsl_name reserves the whole "sarek_" prefix, so a user sarek_fmod is emitted as sarek_sarek_fmod at its definition and at every call.

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 : state -> 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, which is what the emit state's scalar_params field is built from (see state).

Scalar-param shadow renaming

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 reject_float16_kernel : Sarek_ir_types.kernel -> unit
val reject_coopmat_kernel : Sarek_ir_types.kernel -> unit
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.

Omitting it is what every runtime caller does and reproduces the pre-backlog-185 behaviour exactly.

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

Generate complete WGSL source for a kernel.

A special case of generate_with_types with the kernel's OWN type declarations, which is the only thing every production caller ever passed: ~types has exactly the type of the kern_types field (Sarek_ir_types.kernel), so the parameter was redundant with the record it travels in. This used to be a separate 30-80 line copy of the emit sequence that silently omitted record typedefs, variant typedefs and the kernel's variants — source referencing an undeclared struct, with no error. Delegating keeps one emit path per backend.

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).