Module Sarek_matrix_rewrite

The rewritten kernel of a GEMM nest (specs/matrix-pattern-rewrite.md).

A 64 x 64 output tile per block of 128 invocations (4 waves of 32, each a 32 x 32 tile of fragments: 2 x 2 of 16 x 16, or 4 x 4 of 8 x 8 for devices whose matrix unit has 8 x 8 x 8 fragments only, i.e. Metal). Each step of 16 along the reduction stages the A and B tiles from global memory into shared memory — converted to the configuration's operand type, zero outside the region, either layout — then loads fragments from shared memory and multiplies. Every output the region covers is written with the original's initial value and epilogue; nothing else is written.

The rewritten kernel takes the original's parameters, in order, followed by the five int32 parameters of extra_names: the region (Sarek_matrix_mapping.region's me, ne, ze, d, t).

type lowering = Sarek_matrix_scheme.lowering =
  1. | Int32_byte_planes
    (*

    int32: four unsigned byte planes per operand, the ten products of order at most 3, u8*u8+s32->s32: bit-identical modulo 232

    *)
  2. | F16_operands
    (*

    f16*f16+f32->f32 (operands float16, or rounded to it)

    *)
  3. | Bf16_operands
    (*

    float32 operands rounded to bfloat16

    *)
  4. | Bf16x3_operands
    (*

    float32 operands split into three bfloat16 terms, the six products of order at most 2

    *)
val lowering_name : lowering -> string
val config_of : ?frag:int -> lowering -> operand_elt:Sarek_ir_types.elttype -> Sarek_coopmat_types.config

The cooperative-matrix configuration a lowering multiplies with, for the device gate; frag is the fragment edge (m = n = k), 16 by default.

val extra_names : string list
val kernel : ?subgroup:int -> ?frag:int -> Sarek_ir_types.kernel -> Sarek_matrix_pattern.gemm -> lowering -> (Sarek_ir_types.kernel, string) Stdlib.result

Error reason when an operand address, a bound, the initial value or the epilogue cannot be expressed in the rewritten launch, or when subgroup is a width without a template (Sarek_matrix_scheme.supported_subgroup: 32, the default, and 64) or frag is an edge without a template (Sarek_matrix_scheme.supported_fragment: 16, the default, and 8). A block has four waves of subgroup lanes. With 8 the reduction runs as two chained 8-wide multiply-accumulates per step of 16, in k order, into the same accumulator.

val launch_dims : ?subgroup:int -> me:int -> ne:int -> ze:int -> unit -> (int * int * int) * (int * int * int)

(block, grid) as (x, y, z) triples for a region; subgroup as for kernel.