Module Sarek_ir_types

Sarek_ir_types - Pure type definitions for GPU kernel IR

This module contains only type definitions with no external dependencies. Used by spoc_framework for typed generate_source signature.

module Type_id : sig ... end

Runtime type identities with equality proofs.

type memspace =
  1. | Global
  2. | Shared
  3. | Local

Memory spaces

type elttype =
  1. | TInt32
  2. | TInt64
  3. | TFloat16
    (*

    IEEE binary16 storage type. Values are stored/loaded as binary16; arithmetic promotes to TFloat32, computes there, and rounds back on store. There is deliberately no CFloat16 constant: f16 values are produced by conversion (ECast (TFloat16, _)), never by a literal.

    *)
  4. | TFloat32
  5. | TFloat64
  6. | TUint8
    (*

    Unsigned 8-bit integer storage type (backlog-62 slice 3).

    It exists for one reason and its scope is deliberately that narrow: it is the element type of a cooperative-matrix OPERAND BUFFER. Every one of the twelve integer configurations the local RX 7900 XTX advertises has 8-bit operands with a 32-bit accumulator, and coopMatLoad requires the backing array's element type to MATCH the fragment's component type — so there is no route to an integer fragment through a wider buffer, and no way to reach the strict-contract tensor-core path without an 8-bit element type in the IR.

    There is no arithmetic on it. No binop, no literal, no cast to or from it is emitted by any backend; a TUint8 value reaches a kernel only by being read by Sarek_ir_types.coopmat_op.CM_load and leaves only by Sarek_ir_types.coopmat_op.CM_store. That is not an oversight to be filled in later — widening it into a general arithmetic type is a separate decision with its own promotion and overflow questions, and this slice measured none of them.

    *)
  7. | TBool
  8. | TUnit
  9. | TRecord of string * (string * elttype) list
    (*

    Record type: name and field list

    *)
  10. | TVariant of string * (string * elttype list) list
    (*

    Variant type: name and constructor list with arg types

    *)
  11. | TArray of elttype * memspace
    (*

    Array type with element type and memory space

    *)
  12. | TVec of elttype
    (*

    Vector (GPU array parameter)

    *)

Element types

type var = {
  1. var_name : string;
  2. var_id : int;
  3. var_type : elttype;
  4. var_mutable : bool;
}

Variables with type info

type const =
  1. | CInt32 of int32
  2. | CInt64 of int64
  3. | CFloat32 of float
  4. | CFloat64 of float
  5. | CBool of bool
  6. | CUnit

Constants

type binop =
  1. | Add
  2. | Sub
  3. | Mul
  4. | Div
  5. | Mod
  6. | Eq
  7. | Ne
  8. | Lt
  9. | Le
  10. | Gt
  11. | Ge
  12. | And
  13. | Or
  14. | Shl
  15. | Shr
  16. | BitAnd
  17. | BitOr
  18. | BitXor

Binary operators

type unop =
  1. | Neg
  2. | Not
  3. | BitNot

Unary operators

type for_dir =
  1. | Upto
  2. | Downto

Loop direction

type pattern =
  1. | PConstr of string * string list
  2. | PWild

Match pattern

type expr =
  1. | EConst of const
  2. | EVar of var
  3. | EBinop of binop * expr * expr
  4. | EUnop of unop * expr
  5. | EArrayRead of string * expr
    (*

    arridx

    *)
  6. | EArrayReadExpr of expr * expr
    (*

    base_expridx for complex bases

    *)
  7. | ERecordField of expr * string
    (*

    r.field

    *)
  8. | EIntrinsic of string list * string * expr list
    (*

    module path, name, args

    *)
  9. | ECast of elttype * expr
  10. | ETuple of expr list
  11. | EApp of expr * expr list
  12. | ERecord of string * (string * expr) list
    (*

    Record construction: type name, field values

    *)
  13. | EVariant of string * string * expr list
    (*

    Variant construction: type name, constructor, args

    *)
  14. | EArrayLen of string
    (*

    Array length intrinsic

    *)
  15. | EArrayCreate of elttype * expr * memspace
    (*

    elem type, size, memspace

    *)
  16. | EIf of expr * expr * expr
    (*

    condition, then, else - value-returning if

    *)
  17. | EMatch of expr * (pattern * expr) list
    (*

    scrutinee, cases - value-returning match

    *)

Expressions (pure, no side effects)

type lvalue =
  1. | LVar of var
  2. | LArrayElem of string * expr
  3. | LArrayElemExpr of expr * expr
  4. | LRecordField of lvalue * string

L-values (assignable locations)

type stmt =
  1. | SAssign of lvalue * expr
  2. | SSeq of stmt list
  3. | SIf of expr * stmt * stmt option
  4. | SWhile of expr * stmt
  5. | SFor of var * expr * expr * for_dir * stmt
  6. | SMatch of expr * (pattern * stmt) list
  7. | SReturn of expr
  8. | SBarrier
    (*

    Block-level barrier (__syncthreads)

    *)
  9. | SWarpBarrier
    (*

    Warp-level sync (__syncwarp)

    *)
  10. | SExpr of expr
    (*

    Side-effecting expression

    *)
  11. | SEmpty
  12. | SLet of var * expr * stmt
    (*

    Let binding: let v = e in body

    *)
  13. | SLetMut of var * expr * stmt
    (*

    Mutable let: let v = ref e in body

    *)
  14. | SPragma of string list * stmt
    (*

    Pragma hints wrapping a statement

    *)
  15. | SMemFence
    (*

    Memory fence (threadfence)

    *)
  16. | SBlock of stmt
    (*

    Scoped block - creates a C scope for variable isolation

    *)
  17. | SNative of {
    1. gpu : framework:string -> string;
      (*

      Generate GPU code for framework

      *)
    2. ocaml : ocaml_closure;
      (*

      Typed OCaml fallback

      *)
    }
    (*

    Inline native GPU code with OCaml fallback

    *)
  18. | SCoopmat of coopmat_op
    (*

    A cooperative-matrix (tensor-core) operation — backlog-62 slice 3.

    Why ONE statement constructor carrying an operation family, rather than four constructors. Seventeen places in this repository match exhaustively on stmt, and most of them are backends whose only correct response to any of these operations is the same refusal. Four constructors would be sixty-eight arms to write and to keep in agreement; one is seventeen, and a backend that handles SCoopmat at all is then forced by the compiler to consider every member of coopmat_op in one place where the four cases sit next to each other.

    Why fragments are NOT vars and NOT an elttype. A fragment is a subgroup-cooperative value: the whole subgroup collectively holds rows * columns components and each invocation holds a few of them at an implementation-defined position. It cannot be indexed, assigned to, added, cast, passed to a helper, or stored in an array. Giving it an elttype would make every one of those spellable in the IR and would oblige ~36 exhaustive elttype matches to invent an answer for a type none of them can represent. Fragments therefore live in their own namespace, addressed by name, and the only things that can be done to one are the four below.

    *)

Statements (imperative, side effects)

and coopmat_op =
  1. | CM_decl of {
    1. name : string;
    2. frag : Sarek_coopmat_types.fragment;
    }
    (*

    Bring a fragment into scope for the rest of the enclosing block.

    Statement-level rather than a scoping form like SLet, because GLSL, MSL and C all admit a declaration in the middle of a block and because D = A * B + C wants four fragments live at once — nesting four SLet-shaped binders to express that is noise with no invariant behind it.

    *)
  2. | CM_load of {
    1. dst : string;
    2. frag : Sarek_coopmat_types.fragment;
    3. src : string;
    4. index : expr;
    5. stride : expr;
    }
    (*

    Fill dst from the buffer src, row-major, starting at element index, with stride elements between consecutive rows.

    frag is repeated here rather than looked up from the CM_decl: a codegen backend must be able to emit this statement without carrying a fragment environment, and an interpreter must be able to CHECK the two agree. A single source of truth that every consumer has to reconstruct is not a single source of truth.

    Column-major is deliberately absent. It is one more enumerant in GLSL, but it is a second layout to verify on hardware and this slice measured only row-major — an emitted layout nothing has executed is a claim without evidence.

    *)
  3. | CM_store of {
    1. src : string;
    2. frag : Sarek_coopmat_types.fragment;
    3. dst : string;
    4. index : expr;
    5. stride : expr;
    }
    (*

    The inverse of CM_load.

    *)
  4. | CM_muladd of {
    1. dst : string;
    2. a : string;
    3. b : string;
    4. c : string;
    5. cfg : Sarek_coopmat_types.config;
    }
    (*

    dst = a * b + c, the tensor-core instruction itself.

    cfg is the whole point of carrying a configuration rather than four fragments: it is what the device gate is keyed on, it is what says whether the accumulation SATURATES (a property of the operation and not of any operand), and it is what Sarek_coopmat_types.accumulation_is_exact reads to decide whether this statement is under the strict contract or needs the relaxation of docs/design/f16-relaxed-accuracy.md §1.6.

    *)

The four things that can be done with a cooperative-matrix fragment.

Fragment names live in a namespace of their own, separate from var. They are plain strings for the same reason an SShared array name is: a fragment is not an l-value, cannot be captured, and cannot escape the kernel body, so there is nothing for a var's mutability or type field to carry that Sarek_coopmat_types.fragment does not already say.

and decl =
  1. | DParam of var * array_info option
  2. | DLocal of var * expr option
  3. | DShared of string * elttype * expr option

Declarations

and array_info = {
  1. arr_elttype : elttype;
  2. arr_memspace : memspace;
}
and helper_func = {
  1. hf_name : string;
  2. hf_params : var list;
  3. hf_ret_type : elttype;
  4. hf_body : stmt;
}

Helper function (device function called from kernel)

and native_arg =
  1. | NA_Int32 of int32
  2. | NA_Int64 of int64
  3. | NA_Float32 of float
  4. | NA_Float64 of float
  5. | NA_Vec of native_vec

Native argument type for kernel execution. Typed arguments with runtime type witnesses - used by PPX-generated native functions.

and native_vec =
  1. | NV : ('elt, 'underlying) native_vec_ops -> native_vec
and ('elt, 'underlying) native_vec_ops = {
  1. length : int;
  2. elem_size : int;
  3. type_name : string;
  4. type_id : 'elt Type_id.t;
  5. get_f32 : int -> float;
  6. set_f32 : int -> float -> unit;
  7. get_f64 : int -> float;
  8. set_f64 : int -> float -> unit;
  9. get_i32 : int -> int32;
  10. set_i32 : int -> int32 -> unit;
  11. get_i64 : int -> int64;
  12. set_i64 : int -> int64 -> unit;
  13. get_typed : int -> 'elt;
  14. set_typed : int -> 'elt -> unit;
  15. underlying_type_id : 'underlying Type_id.t;
  16. underlying : 'underlying;
}
and ocaml_closure = {
  1. run : block:(int * int * int) -> grid:(int * int * int) -> native_arg array -> unit;
}

Typed Helpers for Custom Types

val vec_get_custom : 'a. 'a Type_id.t -> native_arg -> int -> 'a

Get element from NA_Vec as a type checked custom value.

val vec_set_custom : 'a. 'a Type_id.t -> native_arg -> int -> 'a -> unit

Set element in NA_Vec from a type checked custom value.

val vec_length : native_arg -> int

Get length from NA_Vec

val vec_as_vector : 'a. 'a Type_id.t -> native_arg -> 'a

Get the checked underlying vector/buffer value.

type native_fn_t =
  1. | NativeFn of parallel:bool -> block:(int * int * int) -> grid:(int * int * int) -> native_arg array -> unit

Native function type for V2 execution. Uses typed native_arg.

type kernel = {
  1. kern_name : string;
  2. kern_params : decl list;
  3. kern_locals : decl list;
  4. kern_body : stmt;
  5. kern_types : (string * (string * elttype) list) list;
    (*

    Record type definitions: (type_name, (field_name, field_type); ...)

    *)
  6. kern_variants : (string * (string * elttype list) list) list;
    (*

    Variant type definitions: (type_name, (constructor_name, payload_types); ...)

    *)
  7. kern_funcs : helper_func list;
    (*

    Helper functions defined in kernel scope

    *)
  8. kern_native_fn : native_fn_t option;
    (*

    Optional pre-compiled native function for CPU execution

    *)
}

Kernel representation

val default_kernel : kernel

A kernel with every field at its empty value, for use as the base of a record update: {default_kernel with kern_name = "k"; kern_body = b}.

WHY THIS EXISTS. OCaml requires every field at every record literal, so adding one field to kernel used to mean editing all 119 construction sites in the tree — which is why the type has been avoided rather than extended. A record UPDATE names only the fields it sets, so once a site is written this way a new field costs it nothing.

Prefer this over spelling out the empty fields. make_kernel is the same thing with labels, for new code that would otherwise set most fields.

val make_kernel : ?params:decl list -> ?locals:decl list -> ?types:(string * (string * elttype) list) list -> ?variants:(string * (string * elttype list) list) list -> ?funcs:helper_func list -> ?native_fn:native_fn_t -> name:string -> body:stmt -> unit -> kernel

default_kernel with labels. ~name and ~body are required because a kernel with neither is not a kernel; everything else defaults to empty.