Matrix Units

This page shows how a matrix product (GEMM) in a Sarek program ends up on a GPU’s matrix units, which results that can change, and how to inspect, enable or disable it. Matrix units compute small matrix products in hardware (AMD WMMA on RDNA3, NVIDIA tensor cores); in Sarek’s IR they are cooperative-matrix statements (SCoopmat) of shape 16x16x16 (8x8x8 on Metal). The kernel language has no syntax for them, so a program reaches them in one of three ways:

The int32 rewrite is bit-identical to the original kernel and needs no opt-in, so it is the only one on by default.

Where matrix units run

Sources: Hip_coopmat.ml:43-54 (configurations), Hip_coopmat.ml:57-58 (gfx11 and wave32 rule), Cuda_coopmat.ml (CUDA/C probe; CUDA/PTX keeps coopmat = None in Cuda_plugin_base.ml), Sarek_ir_metal_coopmat.ml (Metal lowering and refusals), Metal_coopmat.ml (probe), Sarek_ir_opencl.ml:616, Sarek_ir_wgsl.ml:1207, Sarek_ir_ptx_stmt.ml:254 (refusals), Sarek_ir_glsl.ml:1395 (Vulkan lowering), Sarek_ir_cuda.ml:574 (CUDA/C), Matrix_specialise.ml:321-333 (Native). The CUDA lowering has been compiled but never run on hardware; the Metal lowering is specified in specs/coopmat-metal-simdgroup.md; the HIP lowering is validated on an RX 7900 XTX (gfx1100) in specs/coopmat-hip-wmma.md. Per-backend defaults of the HIP plugin are in the sarek-hip README.

Compile-time advice

The PPX recognises a GEMM loop nest in a [%kernel ...] and reports it as an OCaml alert named sarek_matrix. It prints at the for loop, does not fail the build, and changes no generated code. For the naive kernel of the matrix multiplication example it prints:

Alert sarek_matrix: Sarek_advice.matrix_product
Sarek: kernel 'matmul_naive' contains a matrix product (GEMM) that matrix units could compute.
  C[m,n] = sum over i = 0..(k - 1) of A[m,i] * B[i,n]
  A = a[((global_idx_y * k) + i)] (row-major, ld k)
  B = b[((i * n) + global_idx_x)] (row-major, ld n)
  C = c[((global_idx_y * n) + global_idx_x)] (row-major, ld n)
  m = global_idx_y, n = global_idx_x
  elements: float32 x float32 -> float32 accumulator
  configuration: 16x16x16 f16*f16+f32->f32 subgroup (operands rounded to float16)
  HIP gfx11 (wave32): executes
  CUDA/C sm_70+: none (the rewrite lowers the integer case only)
  Vulkan: none (the GLSL emitter lowers integer configurations only)
  Metal (Apple GPU): float16 operands, or float32 rounded to float16, onto 8x8x8 fragments (explicit request, or the automatic opt-in); the bfloat16 lowerings: none
  OpenCL, WGSL, PTX, Native: no matrix-unit lowering
  accuracy: relaxed beyond the accumulation order: no matrix unit takes float32 operands, so they are rounded (float16: relative error up to 2^-11 each; bfloat16; or a three-term bfloat16 split, near float32); a rewrite needs Sarek_coopmat.accept_operand_rounding ~precision and accept_relaxed_accumulation
  Execute.run_vectors rewrites it onto matrix units only after the opt-in above AND Sarek_coopmat.accept_automatic_float_rewrite (and neither SAREK_OPT_DISABLE=matrix nor =automatic-float-rewrite); Sarek_matrix.run needs the opt-in above only.
  Silence with [@alert "-sarek_matrix"] on the kernel expression or its binding.

The accuracy line names the opt-ins the kernel would need, and the next line says what Execute.run_vectors does with the nest. The alert is emitted only when the element types have a matrix-unit configuration (see Which kernels are rewritten); a loop that looks like a GEMM but is not (a dot product, a triangular bound, a tiled kernel) gets no advice. A kernel that contains a [%native] block, in its body or a helper, is not analysed. Sources: Sarek_matrix_alert.ml:1-30, :159-160 and :185-189 ([%native]).

Silence it with the standard OCaml alert controls:

open Sarek

let gemm_f32 =
  [%kernel
    fun (a : float32 vector) (b : float32 vector) (c : float32 vector)
        (m : int32) (n : int32) (k : int32) ->
      let open Sarek_stdlib.Std in
      let row = global_idx_y in
      let col = global_idx_x in
      if row < m && col < n then begin
        let sum = mut 0.0 in
        for i = 0 to k - 1l do
          sum := sum +. (a.((row * k) + i) *. b.((i * n) + col))
        done ;
        c.((row * n) + col) <- sum
      end]
[@@alert "-sarek_matrix"]

The attribute also works on the kernel expression, as ([%kernel ...] [@alert "-sarek_matrix"]). For a whole file use [@@@alert "-sarek_matrix"]; for a library, add -alert -sarek_matrix to its (flags ...); to make the advice a build error, use -alert @sarek_matrix.

SAREK_MATRIX_SCAN=<file> makes the PPX append one line per analysed kernel and finding (recognised or near miss) to <file>. It is meant for scanning a code base for GEMM loops. The variable is not a dune dependency, so use a fresh build directory and DUNE_CACHE=disabled, or cached PPX output is not re-run.

Automatic rewrite

Execute.run_vectors asks Matrix_specialise before every launch whether to run a rewritten kernel instead of the one it was given. The rewritten kernel writes exactly the outputs the original launch writes, with one subgroup per output tile instead of one thread per output. Execute.run and raw launches never specialise. Execute.run_vectors_plain launches the kernel as given.

Restricting the rewrite

The environment can only take the rewrite away, never grant it (see Environment variables). SAREK_OPT_DISABLE is read when the program starts:

accumulation-order, operand-rounding-f16 (also -bf16, -bf16x3) revoke the accuracy opt-ins the same way: the program’s call is recorded and has no effect. A restriction is one-way and process-wide, so there is no API to set it in code.

Which kernels are rewritten

A kernel qualifies when all of these hold; otherwise the original runs and the reason is reported (see Reporting).

• The kernel body is one GEMM nest: acc := acc + A[...] * B[...] over a for loop with bounds that are the same for every thread, operands read from global vectors, and a single store to C (optionally with an epilogue such as alpha * acc + beta * C[...]). The kernel writes nothing else (no other array write, no atomics) and has exactly one such nest (only_write_is, Matrix_specialise.ml:108-127).
• Row, column and at most one batch index are affine in the loop variables and thread indices; each operand has a unit stride in one dimension (row-major or column-major). Guards are conjunctions of comparisons with a bound that is the same for every thread.
• No bound, initial value or epilogue reads block_dim_* or grid_dim_*, and no index involves int64.
• The element types have a configuration, below.
• The device is one of HIP with 32-wide waves (gfx11), Metal with 32-wide SIMD-groups (the float16 case only, 8x8x8 fragments), Vulkan (automatically on 64-wide subgroups, the measured width; a 32-wide device only on an explicit request; the block is four subgroups: 128 or 256 threads), or CUDA/C with 32-wide warps (sm_72+, explicit request only); the configuration is offered by the device (Vulkan and CUDA/C: the int32 case only; Metal: the float16 case only); and the rewritten grid fits max_grid_dims.

Sources: element cases Sarek_matrix_advice.ml:43-48; preconditions and the full list of rejected shapes in specs/matrix-pattern-advice.md and specs/matrix-pattern-rewrite.md; bounds, initial value or epilogue reading block_dim_*/grid_dim_* (uses_launch_dims, Sarek_matrix_scheme.ml:294); grid size against max_grid_dims (fits_device, Matrix_specialise.ml:263-296); backend rule Matrix_admission.ml:39-178.

An exception while analysing the kernel gives the original kernel, never a failed launch (Matrix_specialise.ml:150-157 for the analysis, :409-424 for building the rewritten kernel).

Size classes

At each launch the region the original launch covers is computed from the scalar arguments and the grid; the kernel is then one of two shapes (Matrix_specialise.ml:298-318, Sarek_matrix_direct.ml:40-55):

Accuracy opt-ins

A float rewrite changes results, so the program must say it understands how. All opt-ins are process-wide and one-way (there is no way to withdraw one), take a non-empty ~reason (Invalid_argument otherwise), keep the first reason given, and have no environment variable: they are statements in the program’s source (Sarek_coopmat.ml:78-129).

The three calls together enable the automatic float rewrite of a float32 kernel (Sarek_coopmat.mli, section “Automatic float rewrite”):

let () =
  Sarek_coopmat.accept_relaxed_accumulation
    ~reason:"GEMM tolerance checked against the gamma_K bound" ;
  Sarek_coopmat.accept_operand_rounding
    ~precision:Sarek_coopmat.Operands_bf16
    ~reason:"bfloat16 operands are accurate enough for this model" ;
  Sarek_coopmat.accept_automatic_float_rewrite
    ~reason:"let Execute rewrite the model's float GEMMs"

The two accuracy opt-ins alone admit matrix-unit configurations for code that asks for them (Sarek_gemm.Coopmat, Sarek_matrix.run); they do not enable the automatic rewrite. Without accept_automatic_float_rewrite, a float launch through Execute.run_vectors runs the original kernel and the explain line names the missing call. The int32 rewrite needs none of them.

Precision of a float32 product: if the program asks for one (~precision on Sarek_matrix.run or Execute.run_vectors_outcome), that precision must have been accepted. Otherwise float16 is used when accepted, else bfloat16. The bfloat16x3 split reaches near-float32 accuracy but was measured no faster than the scalar kernel, so it is used only on request (Matrix_specialise.ml:206-227). Bounds measured on gfx1100 for each precision are in the Results of specs/matrix-pattern-rewrite.md; they assume finite operands within the narrower type’s range (float16 overflows above 65504).

The bound each configuration is tested against, in units of sum(abs(a*b)) + abs(C) over the reduction (abs(C) only when the product accumulates into a non-zero initial value), with gamma_K = K*u / (1 - K*u), u = 2^-24 and K the number of products added:

The measurements are from specs/matrix-pattern-rewrite.md, Results (RX 7900 XTX, 2026-10-11; shapes 64x64x128, 100x37x129, 256x256x1024; operands uniform in [-1, 1)). A ratio above 1 for the rounded precisions is expected: the operand rounding is a larger error than gamma_K and is covered by the bound of its own row. Only the first row is an accumulation-order error alone; gamma_K is defined in docs/design/f16-relaxed-accuracy.md section 5.2.

The interpreter

By default the interpreter is never rewritten: it stays the independent reference that differential tests compare against, and the rewrite is slower there. It is rewritten only on an explicit request (Sarek_matrix.run, or Execute.run_vectors_outcome ~explicit:true), and not for a launch on a stream. Every other backend is declined before any analysis (Matrix_specialise.ml:321-333).

Reporting and cache

Execute.run_vectors_outcome is run_vectors that also returns the decision:

let launch_and_say ~device ~ir ~args ~block ~grid =
  let outcome =
    Sarek.Execute.run_vectors_outcome ~device ~ir ~args ~block ~grid ()
  in
  print_endline (Sarek_execute.Matrix_specialise.outcome_text outcome)

Every decision about a kernel that is a matrix product is an Sarek_policy.Matrix_rewrite record for Sarek_policy.set_reporter. SAREK_EXPLAIN=1 (exactly 1, read when the program starts) prints each distinct kernel, device and outcome once on stderr, as one line:

sarek-explain <kernel> <framework>:<id> matrix applied proof=<lowering>, <staged|direct (N waves/block)>, <configuration>: bit-identical modulo 2^32
sarek-explain <kernel> <framework>:<id> matrix applied granted=<grants> reason="<the program's reasons>" lowering="<lowering>, <variant>, <configuration>"
sarek-explain <kernel> <framework>:<id> matrix refused <reason>

The first form is the int32 rewrite, the second a float one (the grants it was admitted under: the consent, the accumulation order, the operand rounding). Kernels that are not matrix products and devices that can never rewrite are not reported. Execute.run_vectors_outcome and Sarek_matrix.run return the outcome to the caller as well.

The analysis result and the rewritten kernels are cached per kernel IR (by physical identity), per lowering and size class, so repeated launches of one kernel analyse it once. Restrictions, opt-ins and device are checked at every launch. A kernel whose IR is rebuilt for each launch is analysed each time, as it also misses the launch memo.

Explicit entry points

Sarek_matrix.run takes the device, ir, args, block and grid of Execute.run_vectors plus ?precision, and returns what ran. It also rewrites on the interpreter, and a float kernel needs only the accuracy opt-ins, not accept_automatic_float_rewrite. Execute.run_vectors_outcome ~explicit:true is the same explicit request and waives accept_automatic_float_rewrite in the same way:

let explicit_launch ~device ~ir ~args ~block ~grid =
  match
    Sarek_matrix.run
      ~precision:Sarek_coopmat.Operands_bf16
      ~device ~ir ~args ~block ~grid ()
  with
  | Sarek_matrix.Rewritten _ as outcome ->
      print_endline (Sarek_matrix.outcome_text outcome)
  | Sarek_matrix.Original why -> prerr_endline ("original kernel ran: " ^ why)

Sarek_gemm.Coopmat.run launches a ready-made float16 GEMM D = A x B + C (A m x k and B k x n in float16, C and D m x n in float32, all row-major). The shape must tile exactly (with the default tiling: m and n multiples of 32, k a multiple of 16, and an even number of 32x32 tiles; otherwise Invalid_argument), and it needs accept_relaxed_accumulation and a device with the configuration (Sarek_gemm_coopmat.ml:53-70, :283-309):

let gemm_f16 ~device ~a ~b ~c ~d =
  Sarek_gemm.Coopmat.run ~device ~a ~b ~c ~d ~m:256 ~n:1024 ~k:3072 ()

Related

• Launch window and fusion: the rewritten launch is deferred and fusable like any other.
• Environment variables: every variable on this page.
• Memory and transfers.