Module Sarek_matrix_scheme

module Affine = Sarek_ir_affine
type lowering =
  1. | Int32_byte_planes
  2. | F16_operands
  3. | Bf16_operands
  4. | Bf16x3_operands
val lowering_name : lowering -> string
val bm : int
val bn : int
val elt_bytes : Sarek_ir_types.elttype -> int
val waves : int
val threads_of : subgroup:int -> int
val supported_subgroup : int -> bool
val wave_edge : int
val fragment_edges : int list
val supported_fragment : int -> bool
val fragments_per_wave : frag:int -> int
val extra_names : string list
val i32 : int -> Sarek_ir_types.expr
val gpu : string -> Sarek_ir_types.expr
type fresh = {
  1. mutable next : int;
}
val new_var : fresh -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.var
val max_var_id : Sarek_ir_types.kernel -> int
val let_ : fresh -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.expr -> (Sarek_ir_types.expr -> Sarek_ir_types.stmt) -> Sarek_ir_types.stmt
val for_ : fresh -> string -> upto:int -> (Sarek_ir_types.expr -> Sarek_ir_types.stmt) -> Sarek_ir_types.stmt
val address : Affine.env -> P.operand -> Sarek_ir_types.expr option -> (address, string) Stdlib.result
type scheme = {
  1. planes : int;
  2. shared_elt : Sarek_ir_types.elttype;
  3. operand : C.component_type;
  4. acc : C.component_type;
  5. acc_elt : Sarek_ir_types.elttype;
  6. products : (int * int * int) list;
    (*

    (plane of A, plane of B, accumulator)

    *)
  7. accumulators : int;
  8. stage : Sarek_ir_types.expr -> Sarek_ir_types.expr list;
    (*

    an operand element -> its plane values

    *)
}
val int32_intrinsic : string -> Sarek_ir_types.expr list -> Sarek_ir_types.expr
val pairs : planes:int -> max_order:int -> (int * int * int) list
val scheme : operand_elt:Sarek_ir_types.elttype -> lowering -> scheme
val pad : scheme -> int
val shared_bytes : scheme -> bk:int -> int
val shared_limit : int
val choose_bk : scheme -> int
val config : ?frag:int -> scheme -> C.config
val address_env : Sarek_ir_types.kernel -> P.gemm -> Affine.env
val extra_param : fresh -> string -> Sarek_ir_types.var
val uses_launch_dims : P.gemm -> bool