[Mlir-commits] [mlir] 69ee6b1 - [NVVM] Add FP8 (e4m3/e5m2) support to dense mma.sync (#207307)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 14 00:07:03 PDT 2026
Author: WMC
Date: 2026-07-14T09:06:57+02:00
New Revision: 69ee6b114ac3d3c6652b39d0d57b292bcfa3a81c
URL: https://github.com/llvm/llvm-project/commit/69ee6b114ac3d3c6652b39d0d57b292bcfa3a81c
DIFF: https://github.com/llvm/llvm-project/commit/69ee6b114ac3d3c6652b39d0d57b292bcfa3a81c.diff
LOG: [NVVM] Add FP8 (e4m3/e5m2) support to dense mma.sync (#207307)
## Motivation
I'm adding NVIDIA FP8 mma.sync support in IREE for sm_89/sm_120, and
IREE's codegen lowers through `nvgpu.mma.sync` on this path, which in
turn lowers to `nvvm.mma.sync`. This MLIR-level gap is blocking that
work, so I'm upstreaming the fix here first. This is split into two PRs
per review feedback — this one covers `nvvm.mma.sync` only; a follow-up
PR will add the corresponding `nvgpu.mma.sync` support built on top of
this.
## Description:
`nvvm.mma.sync` (warp-level, dense) don't support FP8 (e4m3/e5m2)
operands, even though the underlying PTX/NVPTX intrinsics have supported
FP8 `mma.sync.m16n8k32` since sm_89. FP8 currently only exists on the
sparse mma path and on warpgroup-level `wgmma`; the warp-level dense
path was never updated.
This adds FP8 (e4m3/e5m2) support to `nvvm.mma.sync`, covering
`m16n8k16` and `m16n8k32`.
## Verification
`$ mlir-opt mlir/test/Dialect/LLVMIR/nvvm.mlir -split-input-file
-verify-diagnostics | FileCheck mlir/test/Dialect/LLVMIR/nvvm.mlir
`
```
%0 = nvvm.mma.sync A[...] B[...] C[...]
{multiplicandAPtxType = #nvvm.mma_type<e4m3>,
multiplicandBPtxType = #nvvm.mma_type<e4m3>,
shape = #nvvm.shape<m = 16, n = 8, k = 32>}
: (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
```
`$ mlir-translate -mlir-to-llvmir mlir/test/Target/LLVMIR/nvvmir.mlir
-split-input-file -verify-diagnostics | FileCheck
mlir/test/Target/LLVMIR/nvvmir.mlir
`
```
%8 = call { float, float, float, float } @llvm.nvvm.mma.m16n8k16.row.col.f32.e4m3.e4m3.f32(i32 %0, i32 %1, i32 %2, float %3, float %4, float %5, float %6)
%11 = call { float, float, float, float } @llvm.nvvm.mma.m16n8k32.row.col.f32.e4m3.e5m2.f32(i32 %0, i32 %1, i32 %2, i32 %3, i32 %4, i32 %5, float %6, float %7, float %8, float %9)
```
All existing tests in nvvm.mlir (including -verify-diagnostics error
cases) still pass — no regression.
Co-authored with Claude Sonnet 5
Signed-off-by: weimin023 <tnwilly at gmail.com>
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
mlir/test/Dialect/LLVMIR/nvvm.mlir
mlir/test/Target/LLVMIR/nvvmir.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 6dd32615451a7..1f76218936e16 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2705,9 +2705,13 @@ class NVVM_MMA_OPS {
list<list<WMMA_REGS>> bit_mma_ops = MMA_OPS<
[GEOM<8,8,128>, GEOM<16,8,128>, GEOM<16,8,256>],
["b1"], [], ["s32"], []>.ret;
+ list<list<WMMA_REGS>> fp8_mma_ops = MMA_OPS<
+ [GEOM<16,8,16>, GEOM<16,8,32>],
+ ["e4m3", "e5m2"], ["e4m3", "e5m2"], ["f16", "f32"], []>.ret;
list<list<WMMA_REGS>> all_mma_sync_ops = !listconcat(
tf32_mma_ops, bf16_mma_ops, f64_mma_ops,
- fp_mma_ops, int_mma_ops, subint_mma_ops, bit_mma_ops);
+ fp_mma_ops, int_mma_ops, subint_mma_ops, bit_mma_ops,
+ fp8_mma_ops);
list<list<WMMA_REGS>> bf16_mma_sp_ops = MMA_OPS<
[GEOM<16,8,16>, GEOM<16,8,32>],
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 14406ab55f138..b29cba96a410f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -1007,6 +1007,15 @@ LogicalResult MmaOp::verify() {
expectedResult.push_back(f16x2x2StructTy);
expectedResult.push_back(f32x4StructTy);
break;
+ case MMATypes::e4m3:
+ case MMATypes::e5m2:
+ // FP8 (m16n8k16 / m16n8k32) packs 4 values per 32-bit register, same
+ // as s8/u8, but the accumulator is f16 or f32 (not integer).
+ kFactor = 16;
+ multiplicandFragType = i32Ty;
+ expectedResult.push_back(f16x2x2StructTy);
+ expectedResult.push_back(f32x4StructTy);
+ break;
case MMATypes::s4:
case MMATypes::u4:
kFactor = 32;
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 72935ffd7b7ce..ae3d18ec46812 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -210,6 +210,54 @@ func.func @nvvm_mma_m16n8k16_bf16_bf16(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i3
llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
}
+// CHECK-LABEL: @nvvm_mma_m16n8k16_e4m3_e4m3
+func.func @nvvm_mma_m16n8k16_e4m3_e4m3(%a0 : i32, %a1 : i32,
+ %b0 : i32,
+ %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) -> !llvm.struct<(f32, f32, f32, f32)> {
+ // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k32_e4m3_e5m2
+func.func @nvvm_mma_m16n8k32_e4m3_e5m2(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,
+ %b0 : i32, %b1 : i32,
+ %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) -> !llvm.struct<(f32, f32, f32, f32)> {
+ // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k16_e4m3_e4m3_f16
+func.func @nvvm_mma_m16n8k16_e4m3_e4m3_f16(%a0 : i32, %a1 : i32,
+ %b0 : i32,
+ %c0 : vector<2xf16>, %c1 : vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)> {
+ // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>, shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k32_e4m3_e5m2_f16
+func.func @nvvm_mma_m16n8k32_e4m3_e5m2_f16(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,
+ %b0 : i32, %b1 : i32,
+ %c0 : vector<2xf16>, %c1 : vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)> {
+ // CHECK: nvvm.mma.sync A[{{.*}}] B[{{.*}}] C[{{.*}}] {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>, multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>, shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+}
+
// CHECK-LABEL: @nvvm_mma_m8n8k16_s8_s8
func.func @nvvm_mma_m8n8k16_s8_s8(%a0 : i32, %b0 : i32,
%c0 : i32, %c1 : i32) -> !llvm.struct<(i32, i32)> {
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index f2888025d8a08..778a5c2229225 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -293,6 +293,54 @@ llvm.func @nvvm_mma_m16n8k16_bf16_bf16(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i3
llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
}
+// CHECK-LABEL: @nvvm_mma_m16n8k16_e4m3_e4m3
+llvm.func @nvvm_mma_m16n8k16_e4m3_e4m3(%a0 : i32, %a1 : i32,
+ %b0 : i32,
+ %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) -> !llvm.struct<(f32, f32, f32, f32)> {
+ // CHECK: call { float, float, float, float } @llvm.nvvm.mma.m16n8k16.row.col.f32.e4m3.e4m3.f32
+ %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1, %c2, %c3]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k32_e4m3_e5m2
+llvm.func @nvvm_mma_m16n8k32_e4m3_e5m2(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,
+ %b0 : i32, %b1 : i32,
+ %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32) -> !llvm.struct<(f32, f32, f32, f32)> {
+ // CHECK: call { float, float, float, float } @llvm.nvvm.mma.m16n8k32.row.col.f32.e4m3.e5m2.f32
+ %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1, %c2, %c3]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
+ llvm.return %0 : !llvm.struct<(f32, f32, f32, f32)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k16_e4m3_e4m3_f16
+llvm.func @nvvm_mma_m16n8k16_e4m3_e4m3_f16(%a0 : i32, %a1 : i32,
+ %b0 : i32,
+ %c0 : vector<2xf16>, %c1 : vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)> {
+ // CHECK: call { <2 x half>, <2 x half> } @llvm.nvvm.mma.m16n8k16.row.col.f16.e4m3.e4m3.f16
+ %0 = nvvm.mma.sync A[%a0, %a1] B[%b0] C[%c0, %c1]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+}
+
+// CHECK-LABEL: @nvvm_mma_m16n8k32_e4m3_e5m2_f16
+llvm.func @nvvm_mma_m16n8k32_e4m3_e5m2_f16(%a0 : i32, %a1 : i32, %a2 : i32, %a3 : i32,
+ %b0 : i32, %b1 : i32,
+ %c0 : vector<2xf16>, %c1 : vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)> {
+ // CHECK: call { <2 x half>, <2 x half> } @llvm.nvvm.mma.m16n8k32.row.col.f16.e4m3.e5m2.f16
+ %0 = nvvm.mma.sync A[%a0, %a1, %a2, %a3] B[%b0, %b1] C[%c0, %c1]
+ {layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>,
+ multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>,
+ shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+ llvm.return %0 : !llvm.struct<(vector<2xf16>, vector<2xf16>)>
+}
+
// f32 return type, f32 accumulate type
// CHECK-LABEL: @nvvm_mma_m16n8k16_f32_f32
llvm.func @nvvm_mma_m16n8k16_f32_f32(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
More information about the Mlir-commits
mailing list