Module Sarek_codegen.Sarek_ir_metal_coopmat

Cooperative matrix (SCoopmat) for the Metal generator, lowered to MSL simdgroup_matrix (specs/coopmat-metal-simdgroup.md).

MSL offers exactly three instantiations, all 8x8 (half, float, bfloat; docs/design/f16-relaxed-accuracy.md §7 slice 6, measured on an Apple M4), and no integer one. The lowered configuration is therefore the single one whose numerics were measured and are a closed form: simdgroup_half8x8 x half8x8 + float8x8 -> float8x8, 8x8x8, bit-equal on 65536/65536 elements to "accumulator = C, then add the eight products in index order, all in binary32". Everything else is refused here, at code generation, with the reason.

Layout is the transpose_matrix flag of simdgroup_load / simdgroup_store: with elements_per_row = stride and no transpose, element (r, c) is at ptr + r * stride + c; with it, at ptr + c * stride + r. This is exactly Sarek_coopmat_types.layout.

module Codegen_error : sig ... end
val unsupported : string -> string -> 'a

What the lowering accepts

val shape_8 : C.shape
val no_integer_reason : string

The reason integer components have no lowering, shared with the TUint8 buffer arm of the type mapper.

val check_fragment : C.fragment -> unit

Refuse, at code generation, a fragment this lowering cannot represent. Total over the fragment type.

val check_config : C.config -> unit

The configurations the lowering emits a multiply-add for.

val buffer_elttype : C.fragment -> Sarek_ir_types.elttype

The kernel buffer element type a fragment loads from or stores to.

val elttype_name : Sarek_ir_types.elttype -> string
val check_buffer : buffers:(string * Sarek_ir_types.elttype) list -> buf:string -> C.fragment -> unit
val check_op : buffers:(string * Sarek_ir_types.elttype) list -> Sarek_ir_types.coopmat_op -> unit
val check_kernel : Sarek_ir_types.kernel -> unit

Whole-kernel check, run before any source is emitted: every fragment, configuration and buffer of every cooperative-matrix statement.

Statements

val msl_elt : C.component_type -> string
val msl_type : C.fragment -> string
val transpose : C.layout -> string
val gen_op : indent:string -> expr_str:(Sarek_ir_types.expr -> string) -> Sarek_ir_types.coopmat_op -> string

gen_op ~indent ~expr_str op is the MSL for one statement. expr_str renders an index or stride expression.

Launch shape

val check_block : Sarek_ir_types.kernel -> (int * int * int) -> unit

A simdgroup_matrix operation needs every lane of its SIMD-group active, so the block must be whole 32-wide SIMD-groups; refused before any compile.