[Mlir-commits] [mlir] [NVGPU] Add FP8 (e4m3/e5m2) support to nvgpu.mma.sync (PR #207342)

[llvmlistbot at llvm.org]

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] [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>
+}