Tensor Core / MMA Builtins

Tensor core builtins implement the Warp Matrix Multiply-Accumulate (WMMA) and Warp Group MMA (WGMMA) interfaces, spanning IDs 678–770 across four SM generations. Each generation added new data types and matrix shapes, resulting in 91 registered builtins that cover half-precision, integer, binary, double-precision, TF32, BF16, and FP8 matrix operations. SM 100 (Blackwell) adds a fifth generation — tcgen05 — documented in Tensor / MMA Codegen.

Key Facts

WMMA Architecture Evolution

HMMA — Half-Precision (IDs 678–707, SM 70+)

The original tensor core builtins provide 16-bit floating-point matrix multiply for three tile shapes. Each shape has 10 operations: load A, load B, load C (f16 and f32 accumulators), store C (f16 and f32), and four MMA variants for input/output precision combinations.

Per-shape operations (10 each):

IMMA — Integer MMA (IDs 708–739, SM 75+)

Integer tensor core operations for INT8 and INT4 data types.

INT8 (IDs 708–731)

Three shapes (16x16x16, 32x8x16, 8x32x16), each with 8 operations:

INT4 (IDs 732–739)

Single shape (8x8x32) with the same operation set but _s4 / _u4 type suffixes.

BMMA — Binary MMA (IDs 740–745, SM 75+)

Binary (1-bit) matrix multiply with XOR-POPC and AND-POPC accumulation modes. Single shape: 8x8x128.

Extended Tensor Core (IDs 746–764, SM 80+)

SM 80 (Ampere) added double-precision, TF32, and BF16 tensor operations.

DMMA — Double Precision (IDs 746, 751–754)

TF32 (IDs 747, 755–757)

BF16 (IDs 748–750, 758–764)

WMMA Lowering Details

Three-Table Lookup

WMMA builtins use a three-table structure for mapping builtin IDs to LLVM intrinsic IDs:

The EDG-side parallel tables live at dword_42810C0 (678–709), dword_4281060 (708–731), dword_4281020 (732–744), addressed from sub_12AC1A0.

Fragment Size Determination

The number of register-level fragments varies by operation and data type:

MMA Codegen Flow

The MMA handler (sub_94E0D0 / sub_12AC5F0) processes 5 input operands:

1. dest_ptr — Pointer to output fragment storage
2. A_fragment — Matrix A input (loaded v100 times)
3. B_fragment — Matrix B input (loaded v95 times)
4. C_fragment — Accumulator input (loaded v101 times)
5. rowcol — Layout operand (validated 0–3 for MMA)

An optional satf flag (saturation, validated 0–1) is consumed for most intrinsics except ID 8279.

The handler emits the MMA call via sub_921880 and scatters results back to the destination fragment through v103 iterations of element-wise stores.

Fragment iteration counts per family (NVVM path, sub_94E0D0):

The output fragment count is determined by bit-test: (0x300C003 >> (intrinsic_id + 127)) & 1 selects 4 vs 8 fragments.

Architecture Gating — Exact Thresholds

The architecture version is stored at *(target_info + 252) as a DWORD.

SM 72 (Xavier) has a unique partial IMMA implementation: only variant 0/1 shapes are supported, with explicit gating that blocks higher variants. This matches hardware reality where Xavier had limited INT8 tensor cores.

WGMMA — Warp Group MMA (SM 90+ Hopper)

WGMMA operates on an entire warp group (4 warps, 128 threads) rather than a single warp. The system is split across four builtin IDs, 20 auxiliary IDs for fence/store/load operations, and two massive handler blocks totaling ~800 lines of lowering logic.

Builtin Registration

Four builtins are registered in sub_90AEE0 (NVVM) and sub_126A910 (EDG):

WGMMA ID Space Overview

The full WGMMA ID range spans 745–770, subdivided into four functional groups:

WGMMA Fence / Commit / Wait (IDs 745–750)

sub_953BA0 (NVVM) / sub_12B1C20 (EDG) builds a red-black tree on first call with 7 entries keyed by builtin ID. Each entry packs:

struct wgmma_fence_entry {
    uint32_t id;           // builtin ID (745–751)
    uint32_t trans_a;      // transpose A flag
    uint32_t shape;        // shape code (0 or 1)
    uint32_t trans_b;      // transpose B flag
    uint32_t a_nregs;      // register count for A fragment
    uint32_t b_nregs;      // register count for B fragment
    uint32_t padding;      // unused alignment
    llvm_type *a_type;     // LLVM type for A (i64, i32, i16x2, i32x4)
    llvm_type *b_type;     // LLVM type for B
    llvm_type *c_type;     // LLVM type for C (i32x2, i32x8)
};

Decoded entries from local variables v47–v106:

Output packed encoding (*a4, 64-bit):

The fence dispatch validates the rowcol operand (must be 0–3) and emits a 4-argument call to intrinsic 9062 (llvm.nvvm.wgmma.fence.aligned) with 3 type overloads. Fragment operands are prepared via sub_94B510.

WGMMA Store (IDs 751–752)

sub_954350 / sub_12B27B0 builds a separate parameter lookup tree. Store operations validate rowcol (0 or 1) and emit a 5-argument call using intrinsic 9145 (llvm.nvvm.wgmma.store) with 2 type overloads. Operands: {constant, B_fragment, descriptor, rowcol, zero}.

WGMMA MMA Async Load (IDs 753–764)

sub_9547E0 (NVVM) / sub_12B2E10 (EDG) builds a 12-entry red-black tree at ctx+656:

Output packed encoding (*a4, 64-bit):

Emits intrinsic 9067 (llvm.nvvm.wgmma.mma.async) with 2 type overloads. Arguments: {constant, B_fragment, rowcol_value, zero_constant}. Results scattered via sub_94B940.

WGMMA MMA Async Compute — The 800-Line Handler (IDs 765–768)

This is the primary WGMMA lowering path. It lives inline in the mega-switch of sub_955A70 (NVVM, lines ~2850–3138) and sub_12B3FD0 (EDG, lines ~2270–3138). The handler implements two completely different intrinsic selection strategies depending on which builtin ID triggered entry.

Argument Extraction

The handler walks the argument chain 7 levels deep from the call expression:

v263 = M dimension              (first constant argument)
v512 = accumulator fragments    (pointer to fragment array)
v528 = A descriptor             (64-bit matrix descriptor or register fragments)
v524 = B descriptor             (64-bit matrix descriptor)
v519 = scale factors            (A and D scale constants)
v264 = layout params            (rowcol encoding)
v516, v265 = shape params       (additional dimension info)
v540 = element type info        (integer type tag from AST)

Each constant argument is validated through sub_620FD0 (EDG) / sub_620FD0 (shared), which extracts the integer value and sets an overflow flag. On overflow:

"unexpected constant overflow in __wgmma_mma_async operand"

This check is applied 5 times: once for N dimension, once for each scale factor, and once for each negate/saturation bit.

Per-Builtin Argument Layouts

Strategy 1: N-Dimension Dispatch (IDs 765–768, inner path)

When the element type is checked and the first argument yields an N dimension, the handler enters a 33-entry switch mapping N values to LLVM intrinsic IDs in the range 10654–10779:

The even/odd intrinsic ID pairing encodes the distinction between integer-element and float-element variants. Type discrimination uses the AST element type: if the element type is integer with width 10 (i.e., a 10-bit integer signaling bf16/tf32 internal encoding), the even (integer) intrinsic is selected; otherwise the odd (float) intrinsic.

N dimension validation:

if ((N & (N - 1)) != 0)
    error("N only supported for powers of two");

This is applied when the N value does not match any case in the 33-entry switch. The N values 8, 16, 32, 64, 128, 256 are powers of two; the intermediate values (24, 40, 48, ..., 248) are non-power-of-two multiples of 8 that are still valid WGMMA dimensions.

Strategy 2: 5-Dimensional Intrinsic Grid (IDs 753–764 path, shared)

For the full WGMMA async variants (handled through sub_12B2E10), the handler selects from a 144-entry intrinsic table spanning IDs 5304–5447, organized as a 5-dimensional grid:

Base addresses and stride:

Size-based opcode selection (for f16, ID 765):

The mapping formula: base + N_offset + shared_offset + type_offset + variant_offset. The accumulator size is extracted by sub_12A71A0(expr) from the expression type chain.

WGMMA Config Bit Packing

Multiple boolean arguments are packed into a single configuration word passed to the final intrinsic call:

Combined via: v79 = bit0 | (bit1 << 1) | (bit2 << 2) | (bit4 << 4).

After intrinsic selection, the handler:

1. Converts the accumulator pointer to a vector pointer (.asvecptr tag)
2. Extracts bitfield from constant operands for mode flags
3. Calls sub_1285290 / sub_921880 with name hint "mmafrag"
4. Scatters results via sub_94B940 / sub_1280F50 (size 4 = float elements)

WGMMA Validation Summary

All constant arguments pass through sub_620FD0, which extracts the integer value and sets an overflow flag.

WGMMA Support Functions

Packed MMA Descriptor Word

The MMA PTX string builder at sub_21E74C0 (AsmPrinter) / sub_35F_range (NVPTX backend) reads a packed 64-bit descriptor for all MMA instruction emission. The descriptor is stored at:

v22 = *(QWORD *)(*(QWORD *)(a1 + 16) + 16 * a2 + 8)

Shape Enum

Unknown shape codes hit the default branch and abort via BUG(). String emission uses fast-path integer stores: *(QWORD *)ptr = 0x36316B386E36316DLL emits "m16n8k16" as a single 8-byte write.

Type Enum

Any other type code produces fatal error: "Wrong MMA element type".

Shape x Type x Architecture Summary

tcgen05 Blackwell Overview (SM 100+)

Full tcgen05 documentation lives in Tensor / MMA Codegen. Key points summarized here for cross-reference:

Data type kinds (bits [8:6] of the tcgen05 operand, emitted by sub_35F3330):

Modifier fields:

Collector modes (emitted by sub_35F38B0):

tcgen05 scaled MMA operand builder (sub_21E8CD0 / sub_35F3E90):

Note the asymmetry: scaleD/transA/transB emit boolean "0"/"1" strings, while negA/negB emit sign multiplier "1"/"-1" strings. This reflects the PTX encoding where negation is a multiplication factor and transpose is a boolean flag.

LLVM Intrinsic Reference

NVPTX Backend Duplicate Functions

All MMA emission functions exist in two structurally identical copies:

The pairs differ only in error reporting (sub_16BD130 vs sub_C64ED0) and reference counting functions (sub_1623A60/sub_161E7C0 vs sub_B96E90/sub_B91220).

Cross-References

• Tensor / MMA Codegen — backend PTX emission, tcgen05 full detail
• NVPTX Opcodes — ISD opcode numbers
• SM 90 (Hopper) — WGMMA architecture context, TMA, cluster
• SM 100 (Blackwell) — tcgen05 architecture context
• Builtin System — hash table, registration, dispatch architecture