Tensor core / MMA intrinsics

[tetsuo-cpp]

Summary

Add support for tensor core matrix multiply-accumulate (MMA) operations on modern NVIDIA GPUs.

Motivation

Tensor cores provide massive throughput for matrix operations (e.g., 16x16x16 matrix multiply in a single instruction). This is essential for competitive matmul performance on modern GPUs (Volta and later).

Design considerations

• WMMA (Warp Matrix Multiply Accumulate) API operates on matrix fragments
• Fragment types: a (M×K), b (K×N), c/d (M×N)
• Supported shapes: 16x16x16, 32x8x16, 8x32x16 (varies by GPU architecture)
• Need words for: load fragment, store fragment, MMA compute

Implementation notes

• Maps to nvvm.wmma or nvvm.mma intrinsics
• Requires floating-point support (depends on f16/f32 types)
• Fragment storage is distributed across warp lanes
• This is a significant feature that may need its own design document

Priority

Nice to have — needed for peak matmul performance but not for correctness.

[tetsuo-cpp]

Tile Stack Design Proposal

Following the Forth tradition of separate stacks (data stack, return stack, float stack), tensor core operations would use a tile stack that holds opaque matrix fragments.

Words

Word	Data Stack	Tile Stack	Description
TZERO	( -- )	( -- acc )	Push zeroed accumulator
TLOAD-A	( addr stride -- )	( -- A )	Load A-operand fragment
TLOAD-B	( addr stride -- )	( -- B )	Load B-operand fragment
TLOAD-C	( addr stride -- )	( -- C )	Load accumulator from memory
TMMA	( -- )	( acc A B -- acc' )	acc' = acc + A*B
TSTORE	( addr stride -- )	( acc -- )	Store accumulator to memory

Addresses and strides live on the regular data stack. Fragments live on the tile stack. The two stacks interact only through load/store words.

Example: GEMM inner loop

\ C += A * B, tiled over K dimension
TZERO                                     \ tile: [acc]
k-tiles BEGIN
  i TILE-SIZE * a-base + a-stride TLOAD-A \ tile: [acc, A_i]
  i TILE-SIZE * b-base + b-stride TLOAD-B \ tile: [acc, A_i, B_i]
  TMMA                                    \ tile: [acc']
LOOP
c-addr c-stride TSTORE                    \ tile: []

MLIR mapping

Each tile stack entry becomes an SSA value of type !gpu.mma_matrix:

• TLOAD-A → gpu.subgroup_mma_load_matrix with "AOp" type
• TLOAD-B → gpu.subgroup_mma_load_matrix with "BOp" type
• TLOAD-C → gpu.subgroup_mma_load_matrix with "COp" type
• TMMA → gpu.subgroup_mma_compute
• TSTORE → gpu.subgroup_mma_store_matrix

The typed load words (TLOAD-A vs TLOAD-B vs TLOAD-C) provide the fragment type annotation the hardware needs, without exposing a type system to the Forth programmer.

Why this works

• Minimal — only ~6 new words
• Forth-idiomatic — separate stacks are standard Forth practice
• Maps to hardware — each word is one WMMA/MMA instruction
• Composable — multiple accumulators are just more entries on the tile stack
• Tile size can be a compile-time parameter or implied by the target architecture (16x16x16 on Volta/Ampere, etc.)