diff --git a/docs/spec/code-organization.md b/docs/spec/code-organization.md index 260c90cdb..3e692cc88 100644 --- a/docs/spec/code-organization.md +++ b/docs/spec/code-organization.md @@ -154,8 +154,8 @@ compilation target nests as `ir/{dialect}/{target}/{category}/.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 diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index 6790e129a..c46081e35 100644 --- a/docs/spec/runtime.md +++ b/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. diff --git a/docs/spec/tir.md b/docs/spec/tir.md index 80a0ec15e..b2daf0dcf 100644 --- a/docs/spec/tir.md +++ b/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.*`) diff --git a/include/tilefoundry/runtime/cuda/ops/mma.cuh b/include/tilefoundry/runtime/cuda/ops/mma.cuh index 6893869d9..51e5af557 100644 --- a/include/tilefoundry/runtime/cuda/ops/mma.cuh +++ b/include/tilefoundry/runtime/cuda/ops/mma.cuh @@ -9,6 +9,7 @@ template __device__ void mma(TA const &a, TB const &b, TC &c) { static_assert( mma_impl::atom_shaped_v, - "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); } diff --git a/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h b/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h index 343ed54fe..5770dbf5a 100644 --- a/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/mma/mma_impl.h @@ -14,8 +14,40 @@ template __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 __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 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 +using element_t = typename tilefoundry::local_view_t::value_type; } @@ -25,20 +57,37 @@ namespace mma_impl { struct Atom { template __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>::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, __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 -inline constexpr bool atom_shaped_v = - frag_v, 8> && - frag_v, 4> && - frag_v, 4>; +inline constexpr bool atom_shaped_v = [] { + using Element = mma_detail::element_t; + using For = mma_detail::AtomFor; + if constexpr (!For::known) + return false; + else + return cute::is_same_v, Element> && + cute::is_same_v, float> && + frag_v, For::a> && + frag_v, For::b> && + frag_v, For::c>; +}(); } diff --git a/include/tilefoundry/runtime/cuda/runtime.cuh b/include/tilefoundry/runtime/cuda/runtime.cuh index 86cee4768..92edb2f88 100644 --- a/include/tilefoundry/runtime/cuda/runtime.cuh +++ b/include/tilefoundry/runtime/cuda/runtime.cuh @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include diff --git a/src/tilefoundry/codegen/cuda/tir/nn/mma.py b/src/tilefoundry/codegen/cuda/tir/nn/mma.py index f03b61ac8..c46004d2d 100644 --- a/src/tilefoundry/codegen/cuda/tir/nn/mma.py +++ b/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}" ) diff --git a/src/tilefoundry/dsl/T/_platforms.py b/src/tilefoundry/dsl/T/_platforms.py index 3eda812f7..a89e10290 100644 --- a/src/tilefoundry/dsl/T/_platforms.py +++ b/src/tilefoundry/dsl/T/_platforms.py @@ -3,6 +3,7 @@ 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 @@ -10,6 +11,10 @@ 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() diff --git a/src/tilefoundry/dsl/_stub_gen.py b/src/tilefoundry/dsl/_stub_gen.py index 57057de44..28a31894e 100644 --- a/src/tilefoundry/dsl/_stub_gen.py +++ b/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", diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma.py b/src/tilefoundry/ir/tir/cuda/nn/mma.py index efa863f00..e4fd49f5c 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma.py +++ b/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", diff --git a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py index 6ed37b36f..f11ff41e4 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py +++ b/src/tilefoundry/ir/tir/cuda/nn/mma_atom.py @@ -12,8 +12,10 @@ LayoutPattern, MeshPattern, Pattern, + ShardLayoutPattern, StarPattern, SwitchPattern, + TensorPattern, WildcardPattern, matched, ) @@ -21,9 +23,10 @@ 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", ] diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py index b9fbae8a2..1475133c7 100644 --- a/src/tilefoundry/ir/tir/cuda/nn/sm80_mma.py +++ b/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"] diff --git a/src/tilefoundry/ir/tir/cuda/nn/sm89_mma.py b/src/tilefoundry/ir/tir/cuda/nn/sm89_mma.py new file mode 100644 index 000000000..9ef8f1055 --- /dev/null +++ b/src/tilefoundry/ir/tir/cuda/nn/sm89_mma.py @@ -0,0 +1,36 @@ +"""The fixed SM89 FP8 warp MMA declaration.""" + +from __future__ import annotations + +from tilefoundry.ir.types import DType, Layout, ShardLayout, Split + +from .mma_atom import MmaAtom, rmem_fragment +from .sm80_mma import WARP +from .sm80_mma import Mma as Sm80Mma + +_A_FRAGMENT = ShardLayout( + layout=Layout(shape=(2, 8, 2, 4, 4), strides=(256, 32, 16, 4, 1)), + attrs=(Split(3), Split(1)), + mesh=WARP, +) +_B_FRAGMENT = ShardLayout( + layout=Layout(shape=(2, 4, 4, 8), strides=(128, 32, 8, 1)), + attrs=(Split(1), Split(3)), + mesh=WARP, +) + + +class Mma(MmaAtom): + """An E4M3 warp MMA, 16 x 8 x 32, accumulating in F32.""" + + namespace = "T.cuda.sm89" + execution_mesh = WARP + capability = "mma.sync" + resource = "tensor_core_engine" + + C = Sm80Mma.C + A = rmem_fragment((16, 32), DType.fp8e4m3, _A_FRAGMENT) + B = rmem_fragment((32, 8), DType.fp8e4m3, _B_FRAGMENT) + + +__all__ = ["Mma"] diff --git a/src/tilefoundry/schedule/facts.py b/src/tilefoundry/schedule/facts.py index 825612540..8df98b8a1 100644 --- a/src/tilefoundry/schedule/facts.py +++ b/src/tilefoundry/schedule/facts.py @@ -14,7 +14,9 @@ def _family_capability( capabilities: tuple[OpCapability, ...], ) -> str | tuple[str, ...] | None: - names = tuple(capability.name for capability in capabilities if capability.name is not None) + names = tuple( + dict.fromkeys(capability.name for capability in capabilities if capability.name is not None) + ) if not names: return None return names[0] if len(names) == 1 else names diff --git a/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.candidates.txt b/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.candidates.txt index ca68f7bb9..f0ba124f9 100644 --- a/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.candidates.txt +++ b/tests/fixtures/schedule/plain/gemm_8192x17408x5120_cta_grid.candidates.txt @@ -48,6 +48,8 @@ target nvidia.h200_sxm mesh: Mesh((Topology("cta", 4352),), layout=(64, 68), names=("bm", "bn")) does not match Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8, 4),), ((*gs, 32, 4, 1),)))) refused T.cuda.sm80.Mma mesh: Mesh((Topology("cta", 4352),), layout=(64, 68), names=("bm", "bn")) does not match Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8),), ((*gs, 1, 4),)))) + refused T.cuda.sm89.Mma + mesh: Mesh((Topology("cta", 4352),), layout=(64, 68), names=("bm", "bn")) does not match Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8),), ((*gs, 1, 4),)))) acc:35 tf.binary per cta lhs=Tensor[(8192, 17408), "f32", ShardLayout( layout=Layout((64, 128, 68, 256), None), diff --git a/tests/fixtures/schedule/tiled_mma.facts.txt b/tests/fixtures/schedule/tiled_mma.facts.txt index 443eadfb4..5b64fb25e 100644 --- a/tests/fixtures/schedule/tiled_mma.facts.txt +++ b/tests/fixtures/schedule/tiled_mma.facts.txt @@ -3,11 +3,13 @@ T.tiled_mma target capability wgmma.mma_async, mma.sync execution mesh atom=T.cuda.sm90.Wgmma Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8, 4),), ((*gs, 32, 4, 1),)))) atom=T.cuda.sm80.Mma Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8),), ((*gs, 1, 4),)))) + atom=T.cuda.sm89.Mma Mesh(('thread',), ComposedLayout(None, p0, Layout(((*g, 4, 8),), ((*gs, 1, 4),)))) predicates: every arrangement has no backward step every arrangement reaches each of its own slots exactly once p0 % 128 == 0 where atom=T.cuda.sm90.Wgmma p0 % 32 == 0 where atom=T.cuda.sm80.Mma + p0 % 32 == 0 where atom=T.cuda.sm89.Mma operands acc the C operand of its atom lhs the A operand of its atom diff --git a/tests/ir/cuda/atoms/_cute_reference.py b/tests/ir/cuda/atoms/_cute_reference.py index 3906f21e9..34060b732 100644 --- a/tests/ir/cuda/atoms/_cute_reference.py +++ b/tests/ir/cuda/atoms/_cute_reference.py @@ -70,7 +70,7 @@ def cute_reference(): assert cuda and (cuda / "cuda_runtime_api.h").is_file(), ( f"Missing CUDA headers: searched CUDA_HOME/include and nvcc/../include ({cuda})" ) - requests = {"mma": "mma bf16 - 0 - 0"} + requests = {"mma-bf16": "mma bf16 - 0 - 0", "mma-fp8e4m3": "mma fp8e4m3 - 0 - 0"} requests.update( { tv_key(dtype, form, n): f"tv {dtype.name} {form.value} {n} K 0" diff --git a/tests/ir/cuda/atoms/cute_reference.cpp b/tests/ir/cuda/atoms/cute_reference.cpp index b6a7584dd..93b9beb90 100644 --- a/tests/ir/cuda/atoms/cute_reference.cpp +++ b/tests/ir/cuda/atoms/cute_reference.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -107,8 +108,10 @@ int main() { std::cout << ','; first = false; std::cout << '"' << key << "\":"; - if (kind == "mma") { + if (kind == "mma" && dtype == "bf16") { emit_tv(); + } else if (kind == "mma" && dtype == "fp8e4m3") { + emit_tv(); } else if (dtype == "bf16") { dispatch(kind, form, n, major, swizzle); } else if (dtype == "f16") { diff --git a/tests/ir/cuda/atoms/test_mma.py b/tests/ir/cuda/atoms/test_mma.py index 3252cc10b..7ddd78128 100644 --- a/tests/ir/cuda/atoms/test_mma.py +++ b/tests/ir/cuda/atoms/test_mma.py @@ -1,4 +1,6 @@ -"""SM80 register fragments compared with CuTe's MMA traits.""" +"""Warp MMA register fragments compared with CuTe's MMA traits.""" + +import pytest from tilefoundry.dsl import T @@ -10,16 +12,21 @@ ) -def test_mma_atom(): +@pytest.mark.parametrize( + ("declaration", "operand_dtype"), + [(T.cuda.sm80.Mma, "bf16"), (T.cuda.sm89.Mma, "fp8e4m3")], + ids=["sm80", "sm89"], +) +def test_mma_atom(declaration, operand_dtype): """Each operand holds the same elements on the same threads as CuTe.""" - reference = cute_reference()["mma"] + reference = cute_reference()[f"mma-{operand_dtype}"] failures = [] for role in ("A", "B", "C"): - atom = T.cuda.sm80.Mma() + atom = declaration() shard = _bound_fragment(atom, role) _check_split_extents_match_mesh(shard) shape = atom.operand_shapes()[("C", "A", "B").index(role)] - _assert_reshard_typeinfer_ok(shape, "f32" if role == "C" else "bf16", shard) + _assert_reshard_typeinfer_ok(shape, "f32" if role == "C" else operand_dtype, shard) ours = to_tv_layout(atom, role) theirs = reference_tv(reference[role]) if ours != theirs: diff --git a/tests/ops/ir/test_matmul.py b/tests/ops/ir/test_matmul.py index f1e2a7b4a..e8b3f6bca 100644 --- a/tests/ops/ir/test_matmul.py +++ b/tests/ops/ir/test_matmul.py @@ -44,6 +44,7 @@ _MM = MatMul() _H200 = CudaTarget("nvidia.h200_sxm") +_SM89 = CudaTarget("nvidia.rtx_6000_ada") _M = Mesh((Topology("gpu", 4),), Layout((4,), (1,)), ("g",)) @@ -282,6 +283,16 @@ def _sharded(shape, attrs): expected=make_tensor_type((64, 32), DType.f32, "rmem"), target=_H200, ), + TypeInferCase( + name="cuda_fp8_registers_read_through_the_sm89_atom", + op=MatMul(out_dtype="f32"), + inputs=( + make_tensor_type((16, 32), DType.fp8e4m3, "rmem"), + make_tensor_type((32, 8), DType.fp8e4m3, "rmem"), + ), + expected=make_tensor_type((16, 8), DType.f32, "rmem"), + target=_SM89, + ), TypeInferCase( name="cuda_on_chip_without_an_mma_is_refused", op=_MM, diff --git a/tests/ops/tir/cuda/test_mma.py b/tests/ops/tir/cuda/test_mma.py index 6193c0f57..b928d1611 100644 --- a/tests/ops/tir/cuda/test_mma.py +++ b/tests/ops/tir/cuda/test_mma.py @@ -151,3 +151,73 @@ def test_fragment_mma_matches_torch_matmul(target: CudaTarget) -> None: torch.cuda.synchronize() expected = a.view(16, 16).float() @ b.view(16, 8).float() assert torch.allclose(c.view(16, 8), expected, rtol=2e-2, atol=2e-2) + + +@module(entry="fp8_host", target=_CUDA, topologies=(Topology("thread", 32),)) +class Fp8MmaTile: + """A 16x32 by 32x8 E4M3 product done as one atom, accumulating in F32.""" + + @prim_func(target=_CUDA) + def fp8_device( + a: Tensor[(16, 32), "fp8e4m3"], + b: Tensor[(32, 8), "fp8e4m3"], + c: Tensor[(16, 8), "f32"], + ): + with Mesh((Topology("thread", 32),), _MESH_LAYOUT, names=("warp", "lane")) as m: + a_frag = T.alloc_tensor( + tensor_type=Tensor[ + (16, 32), + "fp8e4m3", + ((2, 8 @ m.lane, 2, 4 @ m.warp, 4), (256, 32, 16, 4, 1)), + "rmem", + ] + ) + b_frag = T.alloc_tensor( + tensor_type=Tensor[ + (32, 8), + "fp8e4m3", + ((2, 4 @ m.warp, 4, 8 @ m.lane), (128, 32, 8, 1)), + "rmem", + ] + ) + acc = T.alloc_tensor( + tensor_type=Tensor[ + (16, 8), "f32", ((2, 8 @ m.lane, 4 @ m.warp, 2), (64, 8, 2, 1)), "rmem" + ] + ) + a_view = T.tensor_view( + T.ptr_of(a), layout=((2, 8 @ m.lane, 2, 4 @ m.warp, 4), (256, 32, 16, 4, 1)) + ) + b_view = T.tensor_view( + T.ptr_of(b), layout=((2, 4 @ m.warp, 4, 8 @ m.lane), (128, 32, 8, 1)) + ) + T.copy(a_view, a_frag) + T.copy(b_view, b_frag) + T.fill(acc, 0.0) + T.tiled_mma(acc, a_frag, b_frag, atom=T.cuda.sm89.Mma()) + c_view = T.tensor_view( + T.ptr_of(c), layout=((2, 8 @ m.lane, 4 @ m.warp, 2), (64, 8, 2, 1)) + ) + T.copy(acc, c_view) + + @prim_func(target=CpuTarget()) + def fp8_host( + a: Tensor[(16, 32), "fp8e4m3"], + b: Tensor[(32, 8), "fp8e4m3"], + c: Tensor[(16, 8), "f32"], + ): + launch(fp8_device, a, b, c, grid=(1, 1, 1), block=(32, 1, 1)) + + +@pytest.mark.parametrize("target", CUDA_TARGETS) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_fp8_fragment_mma_matches_torch_matmul(target: CudaTarget) -> None: + """E4M3 operands give the F32 product of the same values.""" + torch.manual_seed(1) + a = torch.randn(16, 32, device="cuda").to(torch.float8_e4m3fn) + b = torch.randn(32, 8, device="cuda").to(torch.float8_e4m3fn) + c = torch.empty(16, 8, dtype=torch.float32, device="cuda") + runtime = tilefoundry.compile(retarget(Fp8MmaTile, target), target=target) + runtime(a, b, c) + torch.cuda.synchronize() + torch.testing.assert_close(c, a.float() @ b.float(), rtol=2e-2, atol=2e-2) diff --git a/tests/runtime/cuda/static_asserts.cu b/tests/runtime/cuda/static_asserts.cu index dabceeea9..be8bb20d3 100644 --- a/tests/runtime/cuda/static_asserts.cu +++ b/tests/runtime/cuda/static_asserts.cu @@ -456,3 +456,43 @@ __global__ void k() { static_assert(int(cute::size(cta.layout)) == 32, "unreachable"); } #endif + +#if CASE == 27 +/// mma: an e4m3 A with a bf16 B of the e4m3 atom's length. +__global__ void k(float *p) { + __nv_fp8_e4m3 a[16] = {}; + __nv_bfloat16 b[8] = {}; + float c[4] = {}; + auto at = cute::make_tensor(cute::make_rmem_ptr(&a[0]), cute::Int<16>{}); + auto bt = cute::make_tensor(cute::make_rmem_ptr(&b[0]), cute::Int<8>{}); + auto ct = cute::make_tensor(cute::make_rmem_ptr(&c[0]), cute::Int<4>{}); + mma(at, bt, ct); + p[0] = c[0]; +} +#endif + +#if CASE == 28 +/// mma: e4m3 A and B with a C that is not f32. +__global__ void k(double *p) { + __nv_fp8_e4m3 a[16] = {}, b[8] = {}; + double c[4] = {}; + auto at = cute::make_tensor(cute::make_rmem_ptr(&a[0]), cute::Int<16>{}); + auto bt = cute::make_tensor(cute::make_rmem_ptr(&b[0]), cute::Int<8>{}); + auto ct = cute::make_tensor(cute::make_rmem_ptr(&c[0]), cute::Int<4>{}); + mma(at, bt, ct); + p[0] = c[0]; +} +#endif + +#if CASE == 29 +/// mma: e4m3 operands whose B has the bf16 atom's length. +__global__ void k(float *p) { + __nv_fp8_e4m3 a[16] = {}, b[4] = {}; + float c[4] = {}; + auto at = cute::make_tensor(cute::make_rmem_ptr(&a[0]), cute::Int<16>{}); + auto bt = cute::make_tensor(cute::make_rmem_ptr(&b[0]), cute::Int<4>{}); + auto ct = cute::make_tensor(cute::make_rmem_ptr(&c[0]), cute::Int<4>{}); + mma(at, bt, ct); + p[0] = c[0]; +} +#endif diff --git a/tests/runtime/test_cuda_static_asserts.py b/tests/runtime/test_cuda_static_asserts.py index 2237ea73b..db4622e83 100644 --- a/tests/runtime/test_cuda_static_asserts.py +++ b/tests/runtime/test_cuda_static_asserts.py @@ -37,13 +37,16 @@ 15: "the two projected slices must hold the same number of elements", 16: "ops::dot (warp tier): the fastest axis of the operands' mesh", 17: "ops::dot (block tier): the operands' mesh must be a whole number of warps", - 18: "operands must be the atom's (8, 4, 4) lane fragments", + 18: "operands must be one atom's lane fragments", 19: "ops::sync: a mesh's scope must be cta or thread", 20: "ops::sync: a CTA mesh needs the module's grid-barrier counter", 22: "a reduced mesh axis must divide into whole lanes and whole warps", 24: "ops::copy_async_bulk: both operands must leave the tile whole on every ", 25: "ops::sync: a mesh that skips warps names no barrier", 26: "get: this mesh does not name that level", + 27: "operands must be one atom's lane fragments", + 28: "operands must be one atom's lane fragments", + 29: "operands must be one atom's lane fragments", } pytestmark = pytest.mark.skipif( diff --git a/tests/schedule/test_fixtures.py b/tests/schedule/test_fixtures.py index a393efd59..27d1ed06e 100644 --- a/tests/schedule/test_fixtures.py +++ b/tests/schedule/test_fixtures.py @@ -1876,7 +1876,11 @@ def test_schedule_candidates_reports_every_plain_site( for row in matmul_rows: assert not row["candidates"] - assert {item["id"] for item in row["refused"]} == {"T.cuda.sm90.Wgmma", "T.cuda.sm80.Mma"} + assert {item["id"] for item in row["refused"]} == { + "T.cuda.sm90.Wgmma", + "T.cuda.sm80.Mma", + "T.cuda.sm89.Mma", + } assert all(item["reason"] for item in row["refused"]) if name == "chunk_rmsnorm":