Automerge: [NVGPU] Add FP8 (e4m3/e5m2) support to nvgpu.mma.sync (#207342)

This PR is stacked on #207307

## Description

Add FP8 E4M3/E5M2 support for `nvgpu.mma.sync`, building on the
NVVM-level
support in #207307. This adds NVGPU verification and lowering support,
registers
the supported `m16n8k16`/`m16n8k32` MMA variants, and adds coverage for
NVGPU
lowering and NVVM verification.

## Test

```sh
mlir-opt mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm-fp8.mlir \
  -convert-nvgpu-to-nvvm -split-input-file
```

Co-authored with Claude Sonnet 5

---------

Signed-off-by: weimin023 <tnwilly@gmail.com>

Author: weimin023

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();
 }
 

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();
   }
 

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