Module Sarek_gemm.Coopmat

A tiled f16 GEMM on cooperative-matrix (tensor-core) instructions; see Sarek_gemm_coopmat. Results are bounded, not bit-identical: it needs Sarek_coopmat.accept_relaxed_accumulation before launching.

A tiled f16 GEMM built from cooperative-matrix statements (specs/coopmat-hip-wmma.md).

D = A x B + C for row-major matrices: A is m x k float16, B is k x n float16, C and D are m x n float32. With ~b_layout:Col_major, B is stored transposed (n x k, element (r, c) of B at c * k + r), which is how token-major activations are laid out and what a fragment loads with contiguous reads on HIP gfx11. Each subgroup (wavefront / warp) computes a frag*frags_m x frag*frags_n tile of D (frag is 16 unless given, 8 on Metal), stepping through k frag at a time with one frag x frag x frag multiply-add per fragment pair; a block holds waves subgroups of 32 invocations.

The kernel is IR, not a %kernel: the DSL has no cooperative-matrix surface. It runs where the configuration f16 * f16 + f32 -> f32 is offered — today HIP on gfx11 — and on the interpreter, and in both cases only after Sarek_coopmat.accept_relaxed_accumulation: the matrix unit chooses its addition order, so the result is within gamma_k * (sum |A*B| + |C|) of the exact value (u = 2^-24), not bit-identical to a scalar loop.

type tiling = {
  1. frags_m : int;
  2. frags_n : int;
  3. waves : int;
}
val default_tiling : tiling

2x2 fragments (a 32x32 tile) per wave, 2 waves per block: the best single choice across the spoc_llm prefill shapes measured on the RX 7900 XTX (specs/coopmat-hip-wmma.md, Benchmark). It is not the best per shape: at N = 1024 one fragment per wave fills more of the 96 compute units, and at N = 12288 4x2 is faster.

val subgroup : int
val config_for : frag:int -> Sarek_coopmat.config

The multiply-add configuration for a frag x frag x frag fragment: 16 for the HIP/CUDA/Vulkan matrix units, 8 for Metal's simdgroup_matrix (specs/coopmat-metal-simdgroup.md).

val tile_m : ?frag:int -> tiling -> int
val tile_n : ?frag:int -> tiling -> int
val check_shape : ?frag:int -> tiling -> m:int -> n:int -> k:int -> unit

check_shape ?frag t ~m ~n ~k raises Invalid_argument unless the problem tiles exactly: m a multiple of the wave tile height, n of its width, k of frag (16 unless given), and the number of wave tiles a multiple of t.waves.

val launch_dims : ?frag:int -> tiling -> m:int -> n:int -> Spoc_framework.Framework_sig.dims * Spoc_framework.Framework_sig.dims

(block, grid) for an m x n output; call check_shape first.

The IR

val var : int -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.var
val i32 : int -> Sarek_ir_types.expr
val intrinsic : string -> Sarek_ir_types.expr
val a_name : int -> string
val b_name : int -> string
val acc_name : int -> int -> string
val range : int -> int list
val pairs : tiling -> (int * int) list

One step of the k loop: load the wave's A and B fragments for columns 16*kk .. 16*kk+15 of A, then every multiply-add.

val acc_index : frag:int -> row0:Sarek_ir_types.expr -> col0:Sarek_ir_types.expr -> n:Sarek_ir_types.expr -> (int * int) -> Sarek_ir_types.expr

The accumulator tile at (row0 + 16i, col0 + 16j) in an n-wide buffer.

val wave_body : tiling -> frag:int -> b_layout:Sarek_coopmat.layout -> row0:Sarek_ir_types.expr -> col0:Sarek_ir_types.expr -> k:Sarek_ir_types.expr -> n:Sarek_ir_types.expr -> kk:Sarek_ir_types.var -> Sarek_ir_types.stmt
val ir : ?tiling:tiling -> ?frag:int -> ?b_layout:Sarek_coopmat.layout -> unit -> Sarek_ir_types.kernel

ir ?tiling ?b_layout () is the kernel (a, b, c, d, m, n, k): d = a * b + c. a, b are float16 vectors, c, d float32 vectors, m, n, k int32; b is k x n row-major, or n x k with ~b_layout:Col_major. Launch with launch_dims after check_shape.

Host launch

val ir_cache : (tiling * int * Sarek_coopmat.layout, Sarek_ir_types.kernel) Stdlib.Hashtbl.t
val ir_cache_lock : Stdlib.Mutex.t
val check_lengths : a:('a, 'b) Spoc_core.Vector.t -> b:('c, 'd) Spoc_core.Vector.t -> c:('e, 'f) Spoc_core.Vector.t -> d:('g, 'h) Spoc_core.Vector.t -> m:int -> n:int -> k:int -> unit

run ~device ?tiling ?b_layout ~a ~b ~c ~d ~m ~n ~k () launches ir on device. Raises Invalid_argument on a shape check_shape rejects or a vector shorter than the shape needs, and the launch gate's capability error where the device does not offer the configuration or the program has not called Sarek_coopmat.accept_relaxed_accumulation. Asynchronous like any launch: read d after the usual transfer.

val run : device:Spoc_core.Device.t -> ?tiling:tiling -> ?frag:int -> ?b_layout:Sarek_coopmat.layout -> a:('a, 'b) Spoc_core.Vector.t -> b:('c, 'd) Spoc_core.Vector.t -> c:('e, 'f) Spoc_core.Vector.t -> d:('g, 'h) Spoc_core.Vector.t -> m:int -> n:int -> k:int -> unit -> unit