Module Sarek_codegen.Sarek_ir_cuda_coopmat

Cooperative matrix (SCoopmat) for the CUDA-C generator, which HIP shares.

One generated source serves both toolchains, so every construct here is emitted twice under the platform guard the generator already uses for its f32 barrier: the __HIP_PLATFORM_AMD__ arm is RDNA3 (gfx11) WMMA through the clang builtins, the other arm is nvcuda::wmma. See specs/coopmat-hip-wmma.md for the configurations, the per-lane layout and the reasons for both choices.

What is validated where. The AMD arm is executed on an RX 7900 XTX (gfx1100) against the interpreter: bit-identical for the integer configurations, within the derived bound for f16. The NVIDIA arm is compile-checked only (NVRTC, and ptxas/nvdisasm when present): no NVIDIA GPU was available, and the CUDA backend does not probe coopmat, so the launch gate refuses it before this code is ever compiled for a device.

module Codegen_error : sig ... end
val unsupported : string -> string -> 'a

What the lowering accepts

val shape_16 : C.shape
val is_8bit : C.component_type -> bool
val is_int32 : C.component_type -> bool
val check_fragment : C.fragment -> unit

Refuse, at code generation, a fragment this lowering cannot represent. Total over the fragment type: every case either names its representation or raises with the reason.

val check_config : C.config -> unit

The configurations the lowering emits a multiply-add for.

val buffer_elttype : C.fragment -> Sarek_ir_types.elttype

The kernel buffer element type a fragment may be loaded from or stored to. Integer accumulators of either signedness live in an int32 buffer, and both 8-bit operand types in a uint8 buffer: the signedness is the fragment's, applied by the instruction, exactly as the interpreter narrows the loaded byte.

val elttype_name : Sarek_ir_types.elttype -> string
val check_buffer : buffers:(string * Sarek_ir_types.elttype) list -> buf:string -> C.fragment -> unit

buffers maps a parameter or shared-array name to its element type, when the generator knows it. An unknown name is left to the C++ compiler, which rejects a mismatch too: each helper below takes exactly one pointer type.

val kernel_buffers : Sarek_ir_types.kernel -> (string * Sarek_ir_types.elttype) list

Element types of the kernel's array parameters and shared arrays. A shared array is declared either in kern_locals (DShared) or in the body, as SLet (v, EArrayCreate (elt, _, Shared), _) (how the matrix rewrite declares its staging arrays); private arrays are not buffers.

type operand_layout =
  1. | Loaded of C.layout
  2. | Both_layouts

The layout every CM_load of an operand fragment uses, by fragment name. nvcuda::wmma fixes an operand's layout in its fragment TYPE, so the CUDA arm of the declaration needs it; a name loaded in both layouts has no CUDA spelling. Keyed by name across the whole kernel, scopes ignored: two fragments that share a name in different blocks and are loaded in different layouts get the #error too (conservative, never wrong code).

type context = {
  1. buffers : (string * Sarek_ir_types.elttype) list;
    (*

    element types of the kernel's array parameters and shared arrays

    *)
  2. operand_layouts : (string * operand_layout) list;
}
val operand_layouts : Sarek_ir_types.kernel -> (string * operand_layout) list
val kernel_context : Sarek_ir_types.kernel -> context

What the SCoopmat arm needs from the kernel it is emitting.

Preamble

val preamble : string

The helper definitions, emitted before extern "C" when the kernel has a cooperative-matrix statement: Sarek_ir_cuda_coopmat_preamble.

val relaxed_comment : Sarek_ir_types.kernel -> string

The header comment §6.1 of docs/design/f16-relaxed-accuracy.md requires on a generated source whose results are bounded rather than exact.

Statements

val amd_type : C.fragment -> string
val amd_helper_suffix : C.fragment -> C.layout -> string

Suffix of the AMD helper for a fragment in layout: a_f16, b_i8, acc_f32, acc_cm_f32... An operand helper is named by its access pattern (see the preamble), so a column-major A takes the b helper and a column-major B the a one.

val cuda_elt : C.component_type -> string
val cuda_elt_bytes : C.component_type -> int
val cuda_layout : C.layout -> string
val cuda_type : C.fragment -> C.layout -> string
val both_arms : indent:string -> amd:string -> cuda:string -> string

The two arms of one statement, under the platform guard. Preprocessor directives start at column 0.

val cuda_bf16_unsupported : string
val cuda_arm : C.fragment -> string -> string
val cuda_both_layouts : string -> string
val gen_decl : indent:string -> ctx:context -> string -> C.fragment -> string
val cuda_ptr : const:bool -> C.fragment -> buf:string -> index:string -> string
val cuda_mem_layout : C.layout -> string
val gen_load : indent:string -> buffers:(string * Sarek_ir_types.elttype) list -> expr_str:('a -> string) -> dst:string -> frag:C.fragment -> src:string -> index:'a -> stride:'a -> layout:C.layout -> string
val gen_store : indent:string -> buffers:(string * Sarek_ir_types.elttype) list -> expr_str:('a -> string) -> src:string -> frag:C.fragment -> dst:string -> index:'a -> stride:'a -> layout:C.layout -> string
val gen_muladd : indent:string -> dst:string -> a:string -> b:string -> c:string -> C.config -> string
val gen_op : indent:string -> ctx:context -> expr_str:(Sarek_ir_types.expr -> string) -> Sarek_ir_types.coopmat_op -> string

gen_op ~indent ~ctx ~expr_str op is the source for one cooperative matrix statement, both arms; ctx is kernel_context of the kernel.

Launch shape

val check_block : Sarek_ir_types.kernel -> (int * int * int) -> unit

The lowering runs the matrix instruction on whole wavefronts (wave32 on gfx11) and whole warps; a block whose thread count is not a multiple of 32 has a partial wave whose inactive lanes would feed the instruction undefined operands. Refused before any source is compiled.