Module Sarek_codegen.Sarek_ir_metal

module Codegen_error : sig ... end

Local error module — same raised exception as the package-level Metal_error.

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;
}

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: the field below used to be a module-level ref, so a second generation — on another domain, or simply a later one after Sarek_transpile had written it — read the first one's value.

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.

It is the ONLY field. A mid-refactor draft also carried the framework tag, threaded in as ?framework so a caller could pick the registry spelling; that was dropped once it was clear every caller passed this backend its own name. The record survives the shrink deliberately — a one-field record is the honest shape for "one run's state", and the next thing that needs threading has somewhere to go.

The state for emitting k.

Type Mapping

val mangle_name : string -> string
val metal_type_of_elttype : Sarek_ir_types.elttype -> string

Map Sarek IR element type to Metal C type string

val metal_memspace : Sarek_ir_types.memspace -> string

Map memory space to Metal qualifier

val metal_param_type : Sarek_ir_types.elttype -> string

Map Sarek IR element type to Metal C type for kernel parameters

val metal_helper_param_type : Sarek_ir_types.elttype -> string

Map Sarek IR element type to Metal C type for helper function parameters

val metal_atomic_type_of_elttype : Sarek_ir_types.elttype -> string

Convert type to atomic type for Metal

Thread Intrinsics

val metal_thread_intrinsic : string -> string

Expression Generation

val gen_expr : 'a -> 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_metal_polyfill : 'a -> Stdlib.Buffer.t -> string -> Sarek_ir_types.expr list -> unit

MSL has no cbrt/hypot/expm1/log1p builtins under any name (unlike fabs/rsqrt/atan2, which MSL does define — see Table 6.4 of the Metal Shading Language Specification). These need a multi-token expression instead of a function-name substitution, so they're special-cased here ahead of both the unqualified match arms and the pure registry, applying uniformly to qualified (Float32.cbrt) and unqualified calls alike. cbrt uses sign(x)*pow(abs(x),...) rather than bare pow because pow is undefined for a negative base.

val metal_backend : 'a -> Sarek_ir_types.expr Dispatch.spec

L-value Generation

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

Statement Generation

val indent_nested : string -> string

Nested indentation level

val gen_match_pattern : Stdlib.Buffer.t -> string -> string -> string -> string list -> (string -> Sarek_ir_types.elttype list option) -> unit

Generate match case pattern with variable bindings

val gen_var_decl : 'a -> Stdlib.Buffer.t -> string -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.expr -> unit

Generate variable declaration with initialization

val gen_array_decl : 'a -> Stdlib.Buffer.t -> string -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.expr -> string -> unit

Generate array declaration

val gen_stmt : state -> Stdlib.Buffer.t -> string -> Sarek_ir_types.stmt -> unit

Declaration Generation

val is_vec_type : Sarek_ir_types.elttype -> bool

Check if a type is a vector (requires length parameter). Still used by the Metal-specific gen_param_metal below (buffer-index variant).

val is_pointer_type : Sarek_ir_types.elttype -> bool
val gen_buffer_param : Stdlib.Buffer.t -> string list -> int -> Sarek_ir_types.var -> memspace:Sarek_ir_types.memspace -> elttype:Sarek_ir_types.elttype -> with_length:bool -> int
val gen_param_metal : Stdlib.Buffer.t -> string list -> int -> Sarek_ir_types.decl -> int

Generate parameter with Metal buffer attributes, returns next buffer index

val gen_param : Stdlib.Buffer.t -> Sarek_ir_types.decl -> unit
val collect_atomic_vars_expr : Sarek_ir_types.expr -> string list

Collect variable names used in atomic operations

val collect_atomic_vars_lvalue : Sarek_ir_types.lvalue -> string list
val collect_atomic_vars_stmt : Sarek_ir_types.stmt -> string list
val gen_local : 'a -> Stdlib.Buffer.t -> string -> string list -> Sarek_ir_types.decl -> unit

Helper Function Generation

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

Generate a helper function (Metal device function)

Kernel Generation

val pretty_print_metal : string -> string

Pretty-print Metal source code

val reject_float16_kernel : Sarek_ir_types.kernel -> unit
val reject_float64_kernel : Sarek_ir_types.kernel -> unit
val reject_coopmat_kernel : Sarek_ir_types.kernel -> unit
val metal_fp_contract_pragma : string

The ONLY thing measured to stop Metal contracting a*b+c.

Metal's compile options do NOT do it. Measured on Apple M4 / macOS 15.6.1 (24G90) / Apple clang 17.0.0, on o = a*b + c over 65536 inputs, restricted to the 8773 elements where the DEVICE's own fma differs from the separately-rounded value (so contraction is observable at all):

| build | contracted | |---|---| | default options | 8773 / 8773 | | mathMode = MTLMathModeSafe | **8773 / 8773** | | mathMode=Safe + mathFloatingPointFunctions=Precise | **8773 / 8773** | | fastMathEnabled = NO | **8773 / 8773** | | **this pragma** | **0 / 8773** |

That is §1 corollary 2 again — "a flag that names the hazard is not a mechanism that prevents it" — and it is why the compile options set in Metal_bindings.mtl_compile_options_conformant are NOT a contraction defence and are not described as one. They buy math-function precision; this pragma buys the rounding.

#pragma clang fp contract(off), a volatile thread local, a threadgroup volatile round-trip and an as_type bitcast round-trip were all measured to work too. The pragma is chosen because it is file-scoped, costs no register or memory traffic, and needs no per-expression codegen change — the same reasoning that put precise on GLSL locals (§6). #pragma METAL fp math_mode(safe) does NOT work: like the mathMode property it leaves contraction on. Sweep: tools/probes/metal_contraction_barrier_probe.m.

Sarek's rule is IEEE-754 with every operation rounded as written (docs/fp-contraction-policy.md §1), so this is a conformance requirement, not a tuning choice.

val generate_with_types : types:(string * (string * Sarek_ir_types.elttype) list) list -> Sarek_ir_types.kernel -> string

Generate Metal source with custom type definitions.

Registry lookups are made under the constant tag "Metal". An earlier draft of backlog-185/200 threaded a ?framework argument here so a caller could override it; it was dropped, because every caller passed this backend its own name — exactly the value the constant already is — and the one consumer that would have distinguished them, SNative, now refuses outright.

val generate : Sarek_ir_types.kernel -> string

Generate complete Metal 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.