Module Sarek_codegen.Sarek_ir_ptx_coopmat

Cooperative matrix (SCoopmat) for the PTX generator: the wmma.* instructions of the PTX ISA (tensor cores from sm_70, PTX ISA 6.0).

The lowering accepts exactly what the nvcuda::wmma arm of Sarek_ir_cuda_coopmat accepts: 16x16x16 on subgroup (warp) scope, any layout of A (row/col), B (row/col) and of the accumulator in memory, and two families of configuration.

wmma and not mma.sync: mma.sync.m16n8k16 needs sm_80, and the Turing form m16n8k8 is not translated by ZLUDA (specs/zluda-mma-nvrtc13.md). wmma is the one form sm_70 and later share, and CUDA/C already lowers to it.

Fragments. PTX fixes the per-lane register group of a wmma fragment (m16n16k16): an f16 operand is 8 .b32 registers each holding an f16x2 pair, an 8-bit operand 2 .b32 registers of four bytes, an f32 accumulator 8 .f32 registers and an integer accumulator 8 .b32 registers. The emitter never indexes into a group; it passes the whole group to wmma.load / wmma.mma / wmma.store, which agree on the distribution across lanes.

Precondition, checked at run time. Exactly the CUDA/C one (sarek_cm_require): the element address is 32-byte aligned and stride * element bytes is a multiple of 16, otherwise the kernel executes trap. wmma.load/wmma.store are undefined on a misaligned address; a trap is preferable to wrong results. A .shared array a cooperative-matrix statement accesses is declared 32-byte aligned, so only a misaligned offset into it traps.

What is validated where. Text and refusals: unit tests, any machine. The instructions are checked by ptxas where it is installed. Execution needs an NVIDIA GPU of compute capability 7.0 or later; ZLUDA rejects wmma.*.

val min_sm : int

First SM with wmma, and with 8-bit wmma.

val min_sm_int8 : int
val shape_16 : C.shape

What the lowering accepts

val refuse : string -> string -> 'a
val check_fragment : C.fragment -> unit

Refuse a fragment this lowering cannot represent. Total over the fragment type: every case is accepted or raises with the reason.

val is_8bit : C.component_type -> bool
val is_int32 : C.component_type -> bool
val check_config : C.config -> unit

The configurations a multiply-add is emitted for: f16*f16+f32->f32, and 8-bit operands of one signedness with the same 32-bit integer accumulator type as addend and result (as the CUDA/C arm). The four fragments are checked first.

Layouts

val operand_layouts : Sarek_ir_types.kernel -> (string * C.layout) list

The layout of every operand fragment of the kernel, by name; refuses a name loaded in both layouts. wmma.mma names the layouts of A and B and expects the fragments to have been loaded in them, so a fragment has exactly one (the CUDA/C arm fixes the same thing in the fragment's type).

val accessed_buffers : Sarek_ir_types.kernel -> string list

The arrays a load or store of the kernel accesses. A .shared one among them is declared 32-byte aligned (Sarek_ir_ptx_types.shared_align), so a tile at index 0 passes the address check whatever its neighbours are.

Kernel-level gate

val sm_number : Stdlib.String.t -> int option

sm_number "sm_86" is 86; a suffix (sm_90a) is ignored.

val check_kernel : ?block:(int * int * int) -> sm_target:Stdlib.String.t -> Sarek_ir_types.kernel -> unit

The whole-kernel refusal, run before any instruction is emitted (a per-node refusal would name whichever node the emitter reached first, and a TUint8 buffer that is only stored to reaches none). Also the launch shape: the instructions run on whole warps, and a partial warp has inactive lanes that would feed them undefined operands.

Statements

val regs_text : string list -> string
val layout_suffix : C.layout -> string
val elt_of : C.fragment -> Sarek_ir_types.elttype * int

Buffer element type and byte size a fragment is loaded from or stored to (an integer accumulator of either signedness lives in an int32 buffer, an 8-bit operand of either signedness in a uint8 one).

val elt_name : Sarek_ir_types.elttype -> string
val ptx_type : C.component_type -> string

The PTX type letters of a fragment component in wmma.load/wmma.store (f16, s8, u8, f32, s32).

val frag_reg_count : C.fragment -> int

Registers per lane of a fragment at m16n16k16.

val operand_layout : Sarek_ir_ptx_types.reg_alloc -> string -> C.layout
val emit_decl : Stdlib.Buffer.t -> Sarek_ir_ptx_types.reg_alloc -> string -> C.fragment -> unit
val emit_access : Stdlib.Buffer.t -> Sarek_ir_ptx_types.reg_alloc -> Sarek_ir_ptx_types.env -> what:string -> frag:C.fragment -> buf_name:string -> index:Sarek_ir_types.expr -> stride:Sarek_ir_types.expr -> string * string * string

The element address of a fragment access, after the checks that make the access meaningful: the buffer holds the fragment's element type, and is not per-thread .local memory (a warp-collective access to it has no meaning). Returns the state-space word, the address register and the stride register.

val emit_load : Stdlib.Buffer.t -> Sarek_ir_ptx_types.reg_alloc -> Sarek_ir_ptx_types.env -> dst:string -> frag:C.fragment -> src:string -> index:Sarek_ir_types.expr -> stride:Sarek_ir_types.expr -> layout:C.layout -> unit
val emit_store : Stdlib.Buffer.t -> Sarek_ir_ptx_types.reg_alloc -> Sarek_ir_ptx_types.env -> src:string -> frag:C.fragment -> dst:string -> index:Sarek_ir_types.expr -> stride:Sarek_ir_types.expr -> layout:C.layout -> unit
val emit_muladd : Stdlib.Buffer.t -> Sarek_ir_ptx_types.reg_alloc -> dst:string -> a:string -> b:string -> c:string -> C.config -> unit

emit_op buf alloc env op emits one cooperative-matrix statement.