feat(tir): declare the sm_89 e4m3 warp mma

docs/spec/code-organization.md

@@ -154,8 +154,8 @@ compilation target nests as `ir/{dialect}/{target}/{category}/<name>.py`;
 target-neutral abstractions stay at `ir/{dialect}/{category}/`. For
 example the whole MMA surface is target-owned — `mma.py` defines the `TiledMma`
 op, `mma_atom.py` defines `MmaAtom` / `AtomPattern`, and `sm80_mma.py` /
-`wgmma.py` define the CUDA instruction declarations. All four live under
-`ir/tir/cuda/nn/`. A target-specific rule for a target-neutral Op nests the
+`sm89_mma.py` / `wgmma.py` define the CUDA instruction declarations. All five
+live under `ir/tir/cuda/nn/`. A target-specific rule for a target-neutral Op nests the
 same way: the CUDA `MatMul` type rule lives in `ir/hir/cuda/nn/matmul.py`
 beside no Op class of its own.
 The backend-bound construction stays in TIR: HIR is the checking reference

docs/spec/runtime.md

@@ -1461,13 +1461,17 @@ __device__ void mma(TA const &a, TB const &b, TC &c);
 fragments, and the entry issues one instruction.
 
 - constraints:
-  - `a`, `b`, and `c` contain exactly the atom's `(8, 4, 4)` values per lane.
-    Other shapes are compile-time errors.
+  - The element type of `a` selects the atom: bf16 selects
+    `SM80_16x8x16_F32BF16BF16F32_TN` and e4m3 selects
+    `SM89_16x8x32_F32E4M3E4M3F32_TN`, both accumulating in f32.
+  - `b` holds the element type of `a`, `c` holds f32, and `a`, `b`, and `c`
+    contain exactly that atom's values per lane, `(8, 4, 4)` for bf16 and
+    `(16, 8, 4)` for e4m3. Other shapes and element types are compile-time
+    errors.
   - The fragment map is stated by each operand's `ShardLayout`; the runtime
     publishes no fragment-coordinate function and no accumulator constructor.
-  - Today's atom is `SM80_16x8x16_F32BF16BF16F32_TN`: bf16 operands, f32
-    accumulate. Another instruction is another atom under the same entry, not
-    another entry.
+  - Another instruction is another atom under the same entry, not another
+    entry.
 
 #### 2.6.6 `ops/reduce.cuh`
 
@@ -1641,7 +1645,7 @@ __device__ inline void ldmatrix(Src const &src, Dst &dst);
 **`ldmatrix`.**
 
 One warp loads a dense shared-memory `(16, 16)` bf16 tile into the per-lane A
-fragment consumed by `ops::mma`'s current SM80 atom.
+fragment consumed by `ops::mma`'s SM80 atom.
 
 - constraints:
   - All 32 lanes issue `ldmatrix.sync.aligned.m8n8.x4.shared.b16` together.

docs/spec/tir.md

@@ -1101,7 +1101,9 @@ def variants(
     for an instruction family does not itself state an engine rate.
 
 The public declarations are `T.cuda.sm80.Mma()` (BF16 `16x8x16`, F32
-accumulator, register A/B/C over one warp) and
+accumulator, register A/B/C over one warp), `T.cuda.sm89.Mma()` (E4M3
+`16x8x32`, F32 accumulator laid out as the SM80 one, register A/B/C over one
+warp) and
 `T.cuda.sm90.Wgmma(n=..., dtype=..., form=..., a_major=..., b_major=..., mesh=...)`
 (`dtype` is `bf16`, `f16` or `fp8e4m3`; `64 x n x k` over one warpgroup with
 an F32 accumulator, where `k` is 32 bytes of `dtype`). `dtype` has no default.
@@ -1136,8 +1138,8 @@ T.tiled_mma(acc, lhs, rhs, atom=T.cuda.sm80.Mma())
 
 There is one TIR MMA op: `T.tiled_mma`. There is no optional-atom or bare-MMA
 path. Per-target emission dispatches on the atom ([target](./target.md)); the
-current CUDA emitter accepts the SM80 declaration and rejects WGMMA because
-this stage provides no WGMMA runtime emitter.
+current CUDA emitter accepts the SM80 and SM89 declarations and rejects WGMMA
+because this stage provides no WGMMA runtime emitter.
 
 #### Async copy Ops (`tir.async.*`)
 

include/tilefoundry/runtime/cuda/ops/mma.cuh

@@ -9,6 +9,7 @@ template <class TA, class TB, class TC>
 __device__ void mma(TA const &a, TB const &b, TC &c) {
     static_assert(
         mma_impl::atom_shaped_v<TA, TB, TC>,
-        "ops::mma: operands must be the atom's (8, 4, 4) lane fragments");
+        "ops::mma: operands must be one atom's lane fragments, bf16 (8, 4, 4) "
+        "or e4m3 (16, 8, 4) with f32 C");
     mma_impl::Atom{}(a, b, c);
 }

include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h

@@ -14,8 +14,40 @@ template <class T> __device__ uint16_t as_u16(T const &x) {
     return out;
 }
 
-/// The supported atom.
-using AtomOp = cute::SM80_16x8x16_F32BF16BF16F32_TN;
+/// Four 8-bit values as one register, the first in the low byte.
+__device__ inline uint32_t pack_e4m3x4(uint8_t b0, uint8_t b1, uint8_t b2,
+                                       uint8_t b3) {
+    return uint32_t(b0) | (uint32_t(b1) << 8) | (uint32_t(b2) << 16) |
+           (uint32_t(b3) << 24);
+}
+
+template <class T> __device__ uint8_t as_u8(T const &x) {
+    uint8_t out;
+    __builtin_memcpy(&out, &x, sizeof(uint8_t));
+    return out;
+}
+
+/// The instruction an A element type selects, and the values each lane holds
+/// of A, B and C. An element type with no instruction is not ``known``.
+template <class T> struct AtomFor {
+    static constexpr bool known = false;
+};
+
+template <> struct AtomFor<__nv_bfloat16> {
+    static constexpr bool known = true;
+    using Op = cute::SM80_16x8x16_F32BF16BF16F32_TN;
+    static constexpr int a = 8, b = 4, c = 4;
+};
+
+template <> struct AtomFor<__nv_fp8_e4m3> {
+    static constexpr bool known = true;
+    using Op = cute::SM89_16x8x32_F32E4M3E4M3F32_TN;
+    static constexpr int a = 16, b = 8, c = 4;
+};
+
+/// The element type an operand's lane fragment holds.
+template <class T>
+using element_t = typename tilefoundry::local_view_t<T>::value_type;
 
 }
 
@@ -25,20 +57,37 @@ namespace mma_impl {
 struct Atom {
     template <class TA, class TB, class TC>
     __device__ void operator()(TA const &a, TB const &b, TC &c) const {
-        using CuteAtom = mma_detail::AtomOp;
         using namespace mma_detail;
+        using CuteAtom = typename AtomFor<element_t<TA>>::Op;
 
         auto a_data = a.data();
         auto b_data = b.data();
         auto c_data = c.data();
 
-        uint32_t a0 = pack_bf16x2(as_u16(a_data[0]), as_u16(a_data[1]));
-        uint32_t a1 = pack_bf16x2(as_u16(a_data[4]), as_u16(a_data[5]));
-        uint32_t a2 = pack_bf16x2(as_u16(a_data[2]), as_u16(a_data[3]));
-        uint32_t a3 = pack_bf16x2(as_u16(a_data[6]), as_u16(a_data[7]));
-
-        uint32_t b0 = pack_bf16x2(as_u16(b_data[0]), as_u16(b_data[1]));
-        uint32_t b1 = pack_bf16x2(as_u16(b_data[2]), as_u16(b_data[3]));
+        uint32_t a0, a1, a2, a3, b0, b1;
+        if constexpr (cute::is_same_v<element_t<TA>, __nv_fp8_e4m3>) {
+            a0 = pack_e4m3x4(as_u8(a_data[0]), as_u8(a_data[1]),
+                             as_u8(a_data[2]), as_u8(a_data[3]));
+            a1 = pack_e4m3x4(as_u8(a_data[8]), as_u8(a_data[9]),
+                             as_u8(a_data[10]), as_u8(a_data[11]));
+            a2 = pack_e4m3x4(as_u8(a_data[4]), as_u8(a_data[5]),
+                             as_u8(a_data[6]), as_u8(a_data[7]));
+            a3 = pack_e4m3x4(as_u8(a_data[12]), as_u8(a_data[13]),
+                             as_u8(a_data[14]), as_u8(a_data[15]));
+
+            b0 = pack_e4m3x4(as_u8(b_data[0]), as_u8(b_data[1]),
+                             as_u8(b_data[2]), as_u8(b_data[3]));
+            b1 = pack_e4m3x4(as_u8(b_data[4]), as_u8(b_data[5]),
+                             as_u8(b_data[6]), as_u8(b_data[7]));
+        } else {
+            a0 = pack_bf16x2(as_u16(a_data[0]), as_u16(a_data[1]));
+            a1 = pack_bf16x2(as_u16(a_data[4]), as_u16(a_data[5]));
+            a2 = pack_bf16x2(as_u16(a_data[2]), as_u16(a_data[3]));
+            a3 = pack_bf16x2(as_u16(a_data[6]), as_u16(a_data[7]));
+
+            b0 = pack_bf16x2(as_u16(b_data[0]), as_u16(b_data[1]));
+            b1 = pack_bf16x2(as_u16(b_data[2]), as_u16(b_data[3]));
+        }
 
         float c0 = c_data[0];
         float c1 = c_data[1];
@@ -65,11 +114,21 @@ inline constexpr bool frag_v = [] {
         return int(decltype(cute::size(L{}))::value) == N;
 }();
 
-/// The atom's own operand shapes, stated positively.
+/// The operands of the atom the A element type selects, stated positively: B
+/// holds the same element type as A, C holds f32, and each fragment has the
+/// atom's length.
 template <class TA, class TB, class TC>
-inline constexpr bool atom_shaped_v =
-    frag_v<tilefoundry::local_view_t<TA>, 8> &&
-    frag_v<tilefoundry::local_view_t<TB>, 4> &&
-    frag_v<tilefoundry::local_view_t<TC>, 4>;
+inline constexpr bool atom_shaped_v = [] {
+    using Element = mma_detail::element_t<TA>;
+    using For = mma_detail::AtomFor<Element>;
+    if constexpr (!For::known)
+        return false;
+    else
+        return cute::is_same_v<mma_detail::element_t<TB>, Element> &&
+               cute::is_same_v<mma_detail::element_t<TC>, float> &&
+               frag_v<tilefoundry::local_view_t<TA>, For::a> &&
+               frag_v<tilefoundry::local_view_t<TB>, For::b> &&
+               frag_v<tilefoundry::local_view_t<TC>, For::c>;
+}();
 
 }

include/tilefoundry/runtime/cuda/runtime.cuh

@@ -13,6 +13,7 @@
 #include <cute/atom/mma_atom.hpp>
 #include <cute/atom/mma_traits_sm80.hpp>
 #include <cute/arch/mma_sm80.hpp>
+#include <cute/arch/mma_sm89.hpp>
 #include <cstdint>
 #include <cassert>
 #include <cuda_fp8.h>

src/tilefoundry/codegen/cuda/tir/nn/mma.py

@@ -1,6 +1,7 @@
 """Emit effect-form matrix multiply-accumulate operations.
 
-The SM80 declaration emits one ``tilefoundry::ops::mma`` call. Other
+The SM80 and SM89 declarations each emit one ``tilefoundry::ops::mma`` call,
+which picks the instruction from the operands' element type. Other
 declarations need their own runtime mapping; in particular, WGMMA has no
 emitter in this stage ([runtime §2.6](docs/spec/runtime.md#26-cudaops)).
 """
@@ -10,6 +11,7 @@
 from tilefoundry.ir.core import Var
 from tilefoundry.ir.tir.cuda.nn.mma import TiledMma
 from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma as Sm80Mma
+from tilefoundry.ir.tir.cuda.nn.sm89_mma import Mma as Sm89Mma
 from tilefoundry.target import CudaTarget
 from tilefoundry.visitor_registry.registries import Role, register_codegen
 
@@ -25,7 +27,7 @@ def _emit(call, ctx: CudaCodegenContext) -> None:
     l = ctx.name_for(lhs)
     r = ctx.name_for(rhs)
     atom = call.target.atom
-    if not isinstance(atom, Sm80Mma):
+    if not isinstance(atom, (Sm80Mma, Sm89Mma)):
         raise RuntimeError(
             f"tir.cuda.nn.TiledMma: no codegen handler for atom {atom.reference_name}"
         )

src/tilefoundry/dsl/T/_platforms.py

@@ -3,13 +3,18 @@
 from __future__ import annotations
 
 from tilefoundry.ir.tir.cuda.nn.sm80_mma import Mma
+from tilefoundry.ir.tir.cuda.nn.sm89_mma import Mma as Sm89Mma
 from tilefoundry.ir.tir.cuda.nn.wgmma import Form, Major, Wgmma
 
 
 class _Sm80Namespace:
     Mma = Mma
 
 
+class _Sm89Namespace:
+    Mma = Sm89Mma
+
+
 class _Sm90Namespace:
     Wgmma = Wgmma
     Form = Form
@@ -20,6 +25,7 @@ class _CudaNamespace:
     """``T.cuda`` — compile-time CUDA instruction declarations."""
 
     sm80 = _Sm80Namespace()
+    sm89 = _Sm89Namespace()
     sm90 = _Sm90Namespace()
 
 

src/tilefoundry/dsl/_stub_gen.py

@@ -240,7 +240,7 @@ def constructor(name: str, declaration: type[MmaAtom]) -> str:
         )
 
     namespaces = []
-    for namespace_name in ("sm80", "sm90"):
+    for namespace_name in ("sm80", "sm89", "sm90"):
         namespace = getattr(cuda, namespace_name)
         declarations = sorted(
             (name, value)
@@ -271,6 +271,7 @@ def constructor(name: str, declaration: type[MmaAtom]) -> str:
         [
             "class _Cuda:",
             "    sm80: _CudaSm80",
+            "    sm89: _CudaSm89",
             "    sm90: _CudaSm90",
             "",
             "cuda: _Cuda",

src/tilefoundry/ir/tir/cuda/nn/mma.py

@@ -19,6 +19,7 @@
 
 from .mma_atom import AtomPattern, FromAtom, MmaAtom
 from .sm80_mma import Mma as _Sm80Mma
+from .sm89_mma import Mma as _Sm89Mma
 from .wgmma import Wgmma
 
 _FP_ACC_WIDEN = {
@@ -46,6 +47,11 @@ class TiledMma(Op):
             declaration=_Sm80Mma,
             attribute="atom",
         ),
+        OpCapability(
+            _Sm89Mma.capability,
+            declaration=_Sm89Mma,
+            attribute="atom",
+        ),
     )
 
     @property
@@ -62,7 +68,7 @@ def resource(self):
     atom = ParamDef(
         kind="attribute",
         annotation=MmaAtom,
-        pattern=AtomPattern(Wgmma, _Sm80Mma),
+        pattern=AtomPattern(Wgmma, _Sm80Mma, _Sm89Mma),
     )
     execution_mesh = ParamDef(
         kind="attribute",

src/tilefoundry/ir/tir/cuda/nn/mma_atom.py

@@ -12,18 +12,21 @@
     LayoutPattern,
     MeshPattern,
     Pattern,
+    ShardLayoutPattern,
     StarPattern,
     SwitchPattern,
+    TensorPattern,
     WildcardPattern,
     matched,
 )
 from tilefoundry.ir.pattern import (
     predicates as P,
 )
 from tilefoundry.ir.pattern.utils import declared_shape, selected_pattern
-from tilefoundry.ir.types import ComposedLayout, DType, Layout, Mesh
+from tilefoundry.ir.types import ComposedLayout, DType, Layout, Mesh, ShardLayout
 from tilefoundry.ir.types.layout_algebra import coalesce
 from tilefoundry.ir.types.mesh import levels, starts
+from tilefoundry.ir.types.storage import StorageKind
 
 _MISS = object()
 
@@ -48,6 +51,20 @@ def execution_mesh_pattern(execution_mesh: Mesh, *, groups=False) -> MeshPattern
     return MeshPattern((topology.name,), layout)
 
 
+def rmem_fragment(shape: tuple, dtype: DType, held: ShardLayout) -> TensorPattern:
+    """A warp MMA operand held in registers as *held* lays it out over its mesh."""
+    return TensorPattern(
+        shape=shape,
+        dtype=dtype,
+        storage=StorageKind.RMEM,
+        layout=ShardLayoutPattern(
+            layout=LayoutPattern.from_layout(held.layout),
+            attrs=held.attrs,
+            mesh=execution_mesh_pattern(held.mesh),
+        ),
+    )
+
+
 class MmaAtom(Printable):
     """One instruction declaration; an instance binds its authored parameters."""
 
@@ -242,4 +259,5 @@ def frame(mesh):
     "MmaAtom",
     "physical_frames_match",
     "execution_mesh_pattern",
+    "rmem_fragment",
 ]

src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py

@@ -2,15 +2,9 @@
 
 from __future__ import annotations
 
-from tilefoundry.ir.pattern import (
-    LayoutPattern,
-    ShardLayoutPattern,
-    TensorPattern,
-)
 from tilefoundry.ir.types import DType, Layout, Mesh, ShardLayout, Split, Topology
-from tilefoundry.ir.types.storage import StorageKind as S
 
-from .mma_atom import MmaAtom, execution_mesh_pattern
+from .mma_atom import MmaAtom, rmem_fragment
 
 WARP = Mesh(
     topologies=(Topology("thread", 32),),
@@ -34,18 +28,6 @@
     mesh=WARP,
 )
 
-def _fragment(shape: tuple, dtype, held: ShardLayout) -> TensorPattern:
-    return TensorPattern(
-        shape=shape,
-        dtype=dtype,
-        storage=S.RMEM,
-        layout=ShardLayoutPattern(
-            layout=LayoutPattern.from_layout(held.layout),
-            attrs=held.attrs,
-            mesh=execution_mesh_pattern(WARP),
-        ),
-    )
-
 
 class Mma(MmaAtom):
     """A BF16 warp MMA, 16 x 8 x 16, accumulating in F32."""
@@ -55,9 +37,9 @@ class Mma(MmaAtom):
     capability = "mma.sync"
     resource = "tensor_core_engine"
 
-    C = _fragment((16, 8), DType.f32, _C_FRAGMENT)
-    A = _fragment((16, 16), DType.bf16, _A_FRAGMENT)
-    B = _fragment((16, 8), DType.bf16, _B_FRAGMENT)
+    C = rmem_fragment((16, 8), DType.f32, _C_FRAGMENT)
+    A = rmem_fragment((16, 16), DType.bf16, _A_FRAGMENT)
+    B = rmem_fragment((16, 8), DType.bf16, _B_FRAGMENT)
 
 
 __all__ = ["Mma", "WARP"]