[LLVMGPU] Add NVIDIA FP8 mma.sync intrinsics for sm_89 and sm_120 (#24659)
Part of #24654
## Summary
Adds NVIDIA FP8 `mma.sync` (m16n8k32) strategy-selection support for Ada
(sm_89) and Blackwell (sm_120). IREE previously had no NVIDIA FP8
support at all — FP8 matmuls silently fell back to SIMT tiling on both
targets.
## Not included
FP8 lowering pipeline support. The TileAndFuse LHS transpose is
hardcoded for element[K]=2 (F16); FP8's element[K]=4 causes a
vector.shape_cast shape mismatch. Will do this in the next PR.
## Test (on RTX PRO 6000 Blackwell, sm_120)
```
PATH=/workspace/iree-build/tools:/workspace/iree-build/llvm-project/bin:$PATH \
./llvm-project/bin/llvm-lit -v \
/workspace/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm89.mlir \
/workspace/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir \
/workspace/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
```
```
-- Testing: 3 tests, 3 workers --
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm89.mlir (1 of 3)
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir (2 of 3)
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir (3 of 3)
Testing Time: 0.20s
Total Discovered Tests: 3
Passed: 3 (100.00%)
```
---------
Signed-off-by: weimin023 <tnwilly@gmail.com>
diff --git a/compiler/plugins/target/CUDA/test/target_device_features.mlir b/compiler/plugins/target/CUDA/test/target_device_features.mlir
index 158affd..2fb1072 100644
--- a/compiler/plugins/target/CUDA/test/target_device_features.mlir
+++ b/compiler/plugins/target/CUDA/test/target_device_features.mlir
@@ -14,7 +14,7 @@
// SM89: target_info = #iree_gpu.target<arch = "sm_89", features = "+ptx78",
// SM89-SAME: wgp = <compute = fp64|fp32|fp16|int64|int32|int16|int8, storage = b64|b32|b16|b8,
// SM89-SAME: subgroup = shuffle|arithmetic, dot = dp4xi8toi32,
-// SM89-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
+// SM89-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_MMA_SYNC_F32_16x8x32_F8E4M3FN>, <NV_MMA_SYNC_F32_16x8x32_F8E5M2>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
// SM89-SAME: subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 1024],
// SM89-SAME: max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 101376,
// SM89-SAME: max_workgroup_counts = [2147483647, 65535, 65535]>
@@ -24,7 +24,7 @@
// SM120: target_info = #iree_gpu.target<arch = "sm_120", features = "+ptx87",
// SM120-SAME: wgp = <compute = fp64|fp32|fp16|int64|int32|int16|int8, storage = b64|b32|b16|b8,
// SM120-SAME: subgroup = shuffle|arithmetic, dot = dp4xi8toi32,
-// SM120-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
+// SM120-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_MMA_SYNC_F32_16x8x32_F8E4M3FN>, <NV_MMA_SYNC_F32_16x8x32_F8E5M2>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
// SM120-SAME: subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 64],
// SM120-SAME: max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 101376,
// SM120-SAME: max_workgroup_counts = [2147483647, 65535, 65535]>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
index cc8f1dc..0bc2968 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
@@ -340,6 +340,11 @@
case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16:
case MMAIntrinsic::WMMA_F32_16x16x32_BF16:
return {bf16, bf16, f32};
+ // NVIDIA FP8 mma.sync (m16n8k32, sm_89+).
+ case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN:
+ return {f8E4M3FN, f8E4M3FN, f32};
+ case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2:
+ return {f8E5M2, f8E5M2, f32};
case MMAIntrinsic::WMMAR3_BF16_16x16x16_BF16:
case MMAIntrinsic::WMMAR4_BF16_16x16x16_BF16:
case MMAIntrinsic::WMMA_BF16_16x16x32_BF16:
@@ -872,6 +877,26 @@
default:
return {};
}
+ // FP8 mma.sync (m16n8k32, sm_89+). Same thread arrangement as the f16
+ // m16n8k16 variants; K doubles so element[K] doubles from 2 to 4 (4 FP8
+ // values packed per 32-bit register vs 2 FP16). The accumulator shape is
+ // unchanged (same M×N output). Both supported FP8 types share the same
+ // register layout.
+ case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN:
+ case MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2:
+ switch (operandIndex) {
+ case kMMAOperandLhs: // A [M=16, K=32]: 16 elements/thread
+ return {/*outer=*/{2, 2}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
+ /*element=*/{1, 4}};
+ case kMMAOperandRhs: // B [K=32, N=8]: 8 elements/thread
+ return {/*outer=*/{2, 1}, /*thread=*/{4, 8}, /*strides=*/{1, 4},
+ /*element=*/{4, 1}};
+ case kMMAOperandAcc: // C [M=16, N=8]: 4 elements/thread (same as f16)
+ return {/*outer=*/{2, 1}, /*thread=*/{8, 4}, /*strides=*/{4, 1},
+ /*element=*/{1, 2}};
+ default:
+ return {};
+ }
case MMAIntrinsic::NV_WMMA_F32_16x16x16_F16:
case MMAIntrinsic::NV_WMMA_F16_16x16x16_F16:
return {};
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.td
index 19a64bf..80f6311 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.td
@@ -274,6 +274,10 @@
def NV_MMA_SYNC_F16_16x8x16_F16 : I32EnumAttrCase<"NV_MMA_SYNC_F16_16x8x16_F16", 0x2021>;
def NV_MMA_SYNC_F32_16x8x16_BF16 : I32EnumAttrCase<"NV_MMA_SYNC_F32_16x8x16_BF16", 0x2022>;
+// NV FP8
+def NV_MMA_SYNC_F32_16x8x32_F8E4M3FN : I32EnumAttrCase<"NV_MMA_SYNC_F32_16x8x32_F8E4M3FN", 0x2023>;
+def NV_MMA_SYNC_F32_16x8x32_F8E5M2 : I32EnumAttrCase<"NV_MMA_SYNC_F32_16x8x32_F8E5M2", 0x2024>;
+
// opaque NV intrinsics
def NV_WMMA_F32_16x16x16_F16 : I32EnumAttrCase<"NV_WMMA_F32_16x16x16_F16", 0x2120>;
def NV_WMMA_F16_16x16x16_F16 : I32EnumAttrCase<"NV_WMMA_F16_16x16x16_F16", 0x2121>;
@@ -394,6 +398,9 @@
NV_MMA_SYNC_F32_16x8x16_F16,
NV_MMA_SYNC_F16_16x8x16_F16,
NV_MMA_SYNC_F32_16x8x16_BF16,
+ // NV FP8 intrinsics (sm_89+).
+ NV_MMA_SYNC_F32_16x8x32_F8E4M3FN,
+ NV_MMA_SYNC_F32_16x8x32_F8E5M2,
// opaque NV intrinsics.
NV_WMMA_F32_16x16x16_F16,
NV_WMMA_F16_16x16x16_F16,
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
index 39724bf..65edff2 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
@@ -1020,12 +1020,14 @@
}
const WgpDetails *getAdaWgpDetails() {
- // Ada (sm_89) exposes the same tensor-core instructions as Ampere-class
- // hardware.
+ // Ada (sm_89) adds FP8 mma.sync on top of Ampere-class tensor-core
+ // instructions.
static const MMAIntrinsic mmaOps[] = {
MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16,
MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16,
MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2,
MMAIntrinsic::NV_WMMA_F32_16x16x16_F16,
MMAIntrinsic::NV_WMMA_F16_16x16x16_F16,
};
@@ -1080,6 +1082,8 @@
MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16,
MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16,
MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E4M3FN,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x32_F8E5M2,
MMAIntrinsic::NV_WMMA_F32_16x16x16_F16,
MMAIntrinsic::NV_WMMA_F16_16x16x16_F16,
};
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
index ae991c2..5bc165f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
@@ -16,6 +16,7 @@
name = "lit",
srcs = enforce_glob(
[
+ "config_tile_and_fuse_sm89.mlir",
"config_tile_and_fuse_sm120.mlir",
"config_tile_and_fuse_sm80.mlir",
"config_vector_distribute_sm80.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
index b293f5d..6ca9b6d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
@@ -16,6 +16,7 @@
SRCS
"config_tile_and_fuse_sm120.mlir"
"config_tile_and_fuse_sm80.mlir"
+ "config_tile_and_fuse_sm89.mlir"
"config_vector_distribute_sm80.mlir"
"pipeline_full_smoketests.mlir"
"pipeline_tile_and_fuse_mma_sync.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
index 684c23a..aa6fed1 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
@@ -162,3 +162,63 @@
// CHECK-SAME: reduction = [0, 0, 8]
// CHECK-SAME: subgroup = [2, 4, 0]
// CHECK-SAME: workgroup = [64, 64, 0]
+
+// -----
+
+// Test that FP8 matmul selects NV_MMA_SYNC_F32_16x8x32_F8E4M3FN with TileAndFuse.
+
+#pipeline_layout_fp8 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_f8e4m3fn_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_fp8) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_fp8) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_fp8) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>> -> tensor<256x256xf8E4M3FN>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>> -> tensor<256x256xf8E4M3FN>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf8E4M3FN>, tensor<256x256xf8E4M3FN>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_f8e4m3fn_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x32_F8E4M3FN>
+// CHECK-SAME: promote_operands = [0, 1]
+
+// -----
+
+// Test that F8E5M2 matmul selects NV_MMA_SYNC_F32_16x8x32_F8E5M2.
+
+#pipeline_layout_fp8e5m2 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_f8e5m2_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_fp8e5m2) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_fp8e5m2) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_fp8e5m2) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>> -> tensor<256x256xf8E5M2>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>> -> tensor<256x256xf8E5M2>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf8E5M2>, tensor<256x256xf8E5M2>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_f8e5m2_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x32_F8E5M2>
+// CHECK-SAME: promote_operands = [0, 1]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm89.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm89.mlir
new file mode 100644
index 0000000..4b4fd09
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm89.mlir
@@ -0,0 +1,66 @@
+// RUN: iree-opt --mlir-print-local-scope --split-input-file --iree-gpu-test-target=sm_89 \
+// RUN: --pass-pipeline="builtin.module(iree-llvmgpu-select-lowering-strategy)" %s \
+// RUN: | FileCheck %s
+
+// Tests for Ada (sm_89) FP8 mma.sync strategy selection. F16/BF16 coverage
+// is shared with sm_80 via config_tile_and_fuse_sm80.mlir.
+
+// -----
+
+// Test that F8E4M3FN matmul selects NV_MMA_SYNC_F32_16x8x32_F8E4M3FN.
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_f8e4m3fn_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>> -> tensor<256x256xf8E4M3FN>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E4M3FN>> -> tensor<256x256xf8E4M3FN>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf8E4M3FN>, tensor<256x256xf8E4M3FN>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_f8e4m3fn_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x32_F8E4M3FN>
+// CHECK-SAME: promote_operands = [0, 1]
+
+// -----
+
+// Test that F8E5M2 matmul selects NV_MMA_SYNC_F32_16x8x32_F8E5M2.
+
+#pipeline_layout_e5m2 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_f8e5m2_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_e5m2) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_e5m2) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_e5m2) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>> -> tensor<256x256xf8E5M2>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf8E5M2>> -> tensor<256x256xf8E5M2>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf8E5M2>, tensor<256x256xf8E5M2>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_f8e5m2_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x32_F8E5M2>
+// CHECK-SAME: promote_operands = [0, 1]
