Module Sarek_matrix_mapping

Which outputs a recognised GEMM nest writes (specs/matrix-pattern-rewrite.md, "Preconditions").

A rewritten kernel has its own launch geometry, so it must write exactly the outputs the original launch writes: no more (it would overwrite memory the original leaves alone), no fewer. analyse recognises how the row, column and batch indices derive from the launch, and the guards that bound them; region turns that into numbers for one launch.

type axis =
  1. | X
  2. | Y
  3. | Z
type index_kind =
  1. | Axis of axis
    (*

    global_idx_a, or thread_idx_a + block_dim_a * block_idx_a

    *)
  2. | Loop of {
    1. var : Sarek_ir_types.var;
    2. lo : Sarek_ir_types.expr;
    3. hi : Sarek_ir_types.expr;
    }
    (*

    an enclosing for

    *)
  3. | Quotient of {
    1. linear : string;
    2. divisor : Sarek_ir_types.expr;
    }
    (*

    t / d, t the named linear thread id

    *)
  4. | Remainder of {
    1. linear : string;
    2. divisor : Sarek_ir_types.expr;
    }
    (*

    t mod d

    *)
type role =
  1. | Row
  2. | Col
  3. | Batch
  4. | Linear
type bound =
  1. | Upper of Sarek_ir_types.expr
  2. | Lower of Sarek_ir_types.expr

Upper u: index < u. Lower l: index >= l.

type mapping = {
  1. m : index_kind;
  2. n : index_kind;
  3. batch : index_kind option;
  4. bounds : (role * bound) list;
}
type region = {
  1. me : int;
  2. ne : int;
  3. ze : int;
  4. d : int;
  5. t : int;
}

The outputs of one launch: rows 0, me), columns [0, ne), batches [0, ze); when [d > 0] (a quotient/remainder mapping) also [row * d + column < t].

val axis_char : axis -> char
val analyse : Sarek_ir_types.kernel -> Sarek_matrix_pattern.gemm -> (mapping, string) Stdlib.result

Error reason when the indices or the guards are outside what the rewrite can bound.

val region : mapping -> Sarek_ir_eval_int.launch -> (region, string) Stdlib.result

Error reason when a bound cannot be evaluated, a guard excludes the origin, or the launch writes nothing.