[Mlir-commits] [mlir] [NVGPU] Add FP8 (e4m3/e5m2) support to nvgpu.mma.sync (PR #207342)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 15 09:47:05 PDT 2026
https://github.com/weimin023 updated https://github.com/llvm/llvm-project/pull/207342
>From aca6098aff076bee1214c0a85be73b9648e44711 Mon Sep 17 00:00:00 2001
From: weimin023 <tnwilly at gmail.com>
Date: Fri, 3 Jul 2026 07:43:32 +0000
Subject: [PATCH 1/2] [NVGPU] Add FP8 (e4m3/e5m2) support to nvgpu.mma.sync
Signed-off-by: weimin023 <tnwilly at gmail.com>
---
mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 4 ++++
mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp | 7 ++++---
.../Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir | 12 ++++++++++++
3 files changed, 20 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 948229bb53328..89cc2e5328a84 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -330,6 +330,10 @@ static FailureOr<NVVM::MMATypes> getNvvmMmaType(Type t) {
return NVVM::MMATypes::f64;
if (elType.isF32())
return NVVM::MMATypes::tf32;
+ if (elType.isF8E4M3FN())
+ return NVVM::MMATypes::e4m3;
+ if (elType.isF8E5M2())
+ return NVVM::MMATypes::e5m2;
return failure();
}
diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index c77ecb8431ad8..642e588488641 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -184,7 +184,8 @@ static LogicalResult verifyMmaSyncOp(Operation *op,
numElementA = 1;
numElementB = 1;
} else if (aType.isF32() || aType.isBF16() || aType.isF16() ||
- aType.isInteger(8) || aType.isInteger(4)) {
+ aType.isInteger(8) || aType.isInteger(4) || aType.isF8E4M3FN() ||
+ aType.isF8E5M2()) {
// 8-by-8-128b fundamental tensor core tile size
int operandBitwidth = aType.getIntOrFloatBitWidth();
shapeK = 128 / operandBitwidth; // 128b wide shapeK
@@ -193,8 +194,8 @@ static LogicalResult verifyMmaSyncOp(Operation *op,
numElementB = 32 / operandBitwidth; // 32b wide operand B
} else {
return op->emitError()
- << "expected input data type (i4,i8,f16,bf16,tf32,f64) "
- "supported by "
+ << "expected input data type (i4,i8,f16,bf16,tf32,f64,"
+ "f8E4M3FN,f8E5M2) supported by "
<< op->getName();
}
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
new file mode 100644
index 0000000000000..90c81af872f84
--- /dev/null
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -split-input-file | FileCheck %s
+
+// Test that FP8 (e4m3) nvgpu.mma.sync lowers to nvvm.mma.sync with the
+// correct multiplicand PTX type.
+func.func @fp8_mma_sync(%arg0: vector<4x4xf8E4M3FN>, %arg1: vector<2x4xf8E4M3FN>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
+ // CHECK: nvvm.mma.sync
+ // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>
+ %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<4x4xf8E4M3FN>, vector<2x4xf8E4M3FN>, vector<2x2xf32>) -> vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
>From df07327e29790c7fbf6516637d4f543ed883122f Mon Sep 17 00:00:00 2001
From: weimin023 <tnwilly at gmail.com>
Date: Wed, 15 Jul 2026 16:46:44 +0000
Subject: [PATCH 2/2] [NVGPU] Expand FP8 MMA lowering test coverage
Signed-off-by: weimin023 <tnwilly at gmail.com>
---
.../NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir | 12 -----
.../NVGPUToNVVM/nvgpu-to-nvvm-mma-fp8.mlir | 47 +++++++++++++++++++
2 files changed, 47 insertions(+), 12 deletions(-)
delete mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
create mode 100644 mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-mma-fp8.mlir
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
deleted file mode 100644
index 90c81af872f84..0000000000000
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir
+++ /dev/null
@@ -1,12 +0,0 @@
-// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -split-input-file | FileCheck %s
-
-// Test that FP8 (e4m3) nvgpu.mma.sync lowers to nvvm.mma.sync with the
-// correct multiplicand PTX type.
-func.func @fp8_mma_sync(%arg0: vector<4x4xf8E4M3FN>, %arg1: vector<2x4xf8E4M3FN>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
- // CHECK: nvvm.mma.sync
- // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e4m3>
- // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e4m3>
- // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>
- %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<4x4xf8E4M3FN>, vector<2x4xf8E4M3FN>, vector<2x2xf32>) -> vector<2x2xf32>
- return %0 : vector<2x2xf32>
-}
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-mma-fp8.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-mma-fp8.mlir
new file mode 100644
index 0000000000000..5e7c7f8109f9a
--- /dev/null
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-mma-fp8.mlir
@@ -0,0 +1,47 @@
+// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -split-input-file | FileCheck %s
+
+// CHECK-LABEL: @fp8_mma_e4m3_e4m3_m16n8k16
+func.func @fp8_mma_e4m3_e4m3_m16n8k16(%arg0: vector<2x4xf8E4M3FN>, %arg1: vector<1x4xf8E4M3FN>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
+ // CHECK: nvvm.mma.sync
+ // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>
+ %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<2x4xf8E4M3FN>, vector<1x4xf8E4M3FN>, vector<2x2xf32>) -> vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fp8_mma_e4m3_e4m3_m16n8k32
+func.func @fp8_mma_e4m3_e4m3_m16n8k32(%arg0: vector<4x4xf8E4M3FN>, %arg1: vector<2x4xf8E4M3FN>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
+ // CHECK: nvvm.mma.sync
+ // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e4m3>
+ // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>
+ %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<4x4xf8E4M3FN>, vector<2x4xf8E4M3FN>, vector<2x2xf32>) -> vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fp8_mma_e5m2_e5m2_m16n8k16
+func.func @fp8_mma_e5m2_e5m2_m16n8k16(%arg0: vector<2x4xf8E5M2>, %arg1: vector<1x4xf8E5M2>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
+ // CHECK: nvvm.mma.sync
+ // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e5m2>
+ // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e5m2>
+ // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>
+ %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<2x4xf8E5M2>, vector<1x4xf8E5M2>, vector<2x2xf32>) -> vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fp8_mma_e5m2_e5m2_m16n8k32
+func.func @fp8_mma_e5m2_e5m2_m16n8k32(%arg0: vector<4x4xf8E5M2>, %arg1: vector<2x4xf8E5M2>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
+ // CHECK: nvvm.mma.sync
+ // CHECK-SAME: multiplicandAPtxType = #nvvm.mma_type<e5m2>
+ // CHECK-SAME: multiplicandBPtxType = #nvvm.mma_type<e5m2>
+ // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 32>
+ %0 = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 32]} : (vector<4x4xf8E5M2>, vector<2x4xf8E5M2>, vector<2x2xf32>) -> vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
More information about the Mlir-commits
mailing list