Module Sarek_matrix_pattern

Recognition of matrix-product (GEMM) loop nests in Sarek IR.

A GEMM nest is an accumulator acc (a let mut) updated by a for loop over a reduction variable kv as acc := acc + A[ia] * B[ib] (or fma A[ia] B[ib] acc) and then stored once as C[ic] <- e(acc), where the three indices are affine (Sarek_ir_affine) and separate into a row index M (in A and C), a column index N (in B and C), kv (in A and B) and optional batch indices (in C, and in A or B or both). The full pattern and every rejection are specified in specs/matrix-pattern-advice.md.

This module only recognises. Sarek_matrix_advice says which matrix-unit configuration applies; the PPX turns that into a compile-time alert; and Sarek_fusion uses it to explain why a GEMM operand is not fused.

Scope: the kernel body only, not helper functions. An accumulation written directly into memory (c.(i) <- c.(i) + a * b in the loop) is not a candidate.

include module type of struct include Sarek_matrix_pattern_types end
type layout = Sarek_coopmat_types.layout =
  1. | Row_major
  2. | Col_major

How an operand is laid out in memory, from its two strides: row-major when the column stride is 1, column-major when the row stride is 1. Any other operand (no unit stride, equal strides, overlapping rows) is not a GEMM operand: a matrix unit loads a tile with one base and one leading dimension.

type operand = Sarek_matrix_pattern_types.operand = {
  1. array : string;
  2. elt : Sarek_ir_types.elttype;
  3. index : Sarek_ir_types.expr;
  4. row_stride : Sarek_ir_types.expr;
  5. col_stride : Sarek_ir_types.expr;
  6. layout : layout;
  7. offset : Sarek_ir_types.expr;
}

One matrix of C = A × B. Rows and columns are those of the matrix: A is M × K, B is K × N, C is M × N. The strides are uniform expressions; offset is everything in the index that is neither the row nor the column term (constants and batch terms). index is the let-inlined index.

type update = Sarek_matrix_pattern_types.update =
  1. | Mul_add
    (*

    acc + x * y

    *)
  2. | Fused_multiply_add
    (*

    fma x y acc

    *)
type epilogue = Sarek_matrix_pattern_types.epilogue =
  1. | Plain
    (*

    C[ic] <- acc, possibly through a cast

    *)
  2. | Epilogue of Sarek_ir_types.expr
    (*

    acc occurs once, beside uniform values and reads of C[ic] (e.g. alpha * acc + beta * C[ic])

    *)
type gemm = Sarek_matrix_pattern_types.gemm = {
  1. acc : Sarek_ir_types.var;
  2. init : Sarek_ir_types.expr;
    (*

    the accumulator's initial value, let-inlined

    *)
  3. k_var : Sarek_ir_types.var;
  4. k_lo : Sarek_ir_types.expr;
    (*

    inclusive, uniform

    *)
  5. k_hi : Sarek_ir_types.expr;
    (*

    inclusive, uniform

    *)
  6. m_index : Sarek_ir_types.expr;
    (*

    the row atom, e.g. global_idx_y or tid / n

    *)
  7. n_index : Sarek_ir_types.expr;
    (*

    the column atom

    *)
  8. batch : Sarek_ir_types.expr list;
    (*

    batch atoms of C

    *)
  9. a : operand;
  10. b : operand;
  11. c : operand;
  12. update : update;
  13. epilogue : epilogue;
  14. guards : Sarek_ir_types.expr list;
    (*

    conditions of the enclosing ifs, outermost first, then the store's own guard; let-inlined

    *)
  15. enclosing_loops : Sarek_ir_types.var list;
    (*

    outermost first

    *)
  16. unmodelled_control : string option;
    (*

    control on the path to the nest that guards does not describe (an else branch, a while, a match case): guards is then not the whole condition under which the store runs

    *)
}
type reason = Sarek_matrix_pattern_types.reason =
  1. | Downto_loop
  2. | Loop_body_not_single_update
  3. | Not_a_product
  4. | Operand_not_an_element of Sarek_ir_types.expr
  5. | Operand_not_global of string
  6. | Bounds_not_uniform
  7. | Bounds_data_dependent
  8. | Index_not_affine of string * Sarek_ir_types.expr option
  9. | Int64_arithmetic
    (*

    an index, bound or guard computes in 64 bits (compiler review C-11)

    *)
  10. | Operand_independent_of_reduction of string
  11. | Output_depends_on_reduction
  12. | No_row_column_split
  13. | Stray_index_atom of string
  14. | Nested_accumulation
  15. | Accumulator_escapes
  16. | No_output_store
  17. | Unsupported_epilogue
  18. | Output_aliases_operand of string
  19. | Operand_written of string
  20. | Overlapping_operand of string
  21. | Constant_leading_dimension of string
  22. | Output_written_elsewhere of string
  23. | Operand_not_unit_stride of string
  24. | Dependent_indices
  25. | Unsupported_init

Why a reduction loop is not a GEMM nest. Payloads name the array or the sub-expression at fault.

type near_miss = Sarek_matrix_pattern_types.near_miss = {
  1. nm_acc : Sarek_ir_types.var;
  2. nm_loop : Sarek_ir_types.var;
  3. nm_reason : reason;
}

A reduction loop (a for whose body updates a let mut from itself) that is not a GEMM nest.

type finding = Sarek_matrix_pattern_types.finding =
  1. | Gemm of gemm
  2. | Near_miss of near_miss
val reason_text : reason -> string
val layout_name : layout -> string
val analyze : Sarek_ir_types.kernel -> finding list

Every reduction loop of the kernel body, in body order.

val recognise : Sarek_ir_types.kernel -> gemm list

The Gemm findings of analyze.