[LLVMGPU] Add sm_120 TileAndFuse config coverage, consolidate MMA sync pipeline tests (#24577)
## Description
Adds LLVMGPU TileAndFuse config-selection coverage for the CUDA `sm_120`
target and consolidates MMA sync pipeline-lowering coverage across
NVIDIA
GPU architectures.
The `sm_120` config-selection test verifies that f16 matmul cases select
the
NVIDIA MMA TileAndFuse path before any architecture-specific tuning is
introduced. The pipeline-lowering tests are kept target-independent
because
they start from IR with an attached `mma_sync` intrinsic and only verify
that
the selected lowering pipeline generates the corresponding
`nvgpu.mma.sync`
ops.
Covered `sm_120` config-selection cases:
- `f16xf16 -> f32` matmul selects TileAndFuse with `NV_MMA_SYNC`
- `f16xf16 -> f16` matmul selects TileAndFuse with `NV_MMA_SYNC`
- `f16xf16 -> f32` matmul accumulate sets `convert_acc_gemm`
Covered target-independent pipeline-lowering cases:
- `f16xf16 -> f32` TileAndFuse lowering generates `nvgpu.mma.sync`
- `f16xf16 -> f16` TileAndFuse lowering generates `nvgpu.mma.sync`
- `f16xf16 -> f32` VectorDistribute lowering generates `nvgpu.mma.sync`
- `f16xf16 -> f16` VectorDistribute lowering generates `nvgpu.mma.sync`
This is baseline regression coverage only; it does not introduce new
tile
heuristics or `sm_120`-specific tuning.
## Test
- `llvm-lit -v \
compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
\
compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
\
compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
\
compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_mma_sync.mlir`
```
-- Testing: 4 tests, 4 workers --
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir (1 of 4)
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir (2 of 4)
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_mma_sync.mlir (3 of 4)
PASS: IREE :: src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir (4 of 4)
Testing Time: 0.14s
Total Discovered Tests: 4
Passed: 4 (100.00%)
```
---------
Signed-off-by: weimin023 <tnwilly@gmail.com>
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 07b0ee6..e5ded30 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
@@ -17,11 +17,12 @@
srcs = enforce_glob(
[
"bf16_mma_sm80.mlir",
+ "config_tile_and_fuse_sm120.mlir",
"config_tile_and_fuse_sm80.mlir",
"config_vector_distribute_sm80.mlir",
"pipeline_full_smoketests.mlir",
- "pipeline_tile_and_fuse_sm80.mlir",
- "pipeline_vector_distribute_sm80.mlir",
+ "pipeline_tile_and_fuse_mma_sync.mlir",
+ "pipeline_vector_distribute_mma_sync.mlir",
],
include = ["*.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 1608e8b..5d58dde 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
@@ -15,11 +15,12 @@
lit
SRCS
"bf16_mma_sm80.mlir"
+ "config_tile_and_fuse_sm120.mlir"
"config_tile_and_fuse_sm80.mlir"
"config_vector_distribute_sm80.mlir"
"pipeline_full_smoketests.mlir"
- "pipeline_tile_and_fuse_sm80.mlir"
- "pipeline_vector_distribute_sm80.mlir"
+ "pipeline_tile_and_fuse_mma_sync.mlir"
+ "pipeline_vector_distribute_mma_sync.mlir"
TOOLS
FileCheck
iree-opt
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
new file mode 100644
index 0000000..89cba20
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
@@ -0,0 +1,94 @@
+// RUN: iree-opt --mlir-print-local-scope --split-input-file --iree-gpu-test-target=sm_120 \
+// RUN: --pass-pipeline="builtin.module(iree-llvmgpu-select-lowering-strategy)" %s \
+// RUN: | FileCheck %s
+
+// Test that sm_120 selects #iree_gpu.pipeline<TileAndFuse> with NV_MMA_SYNC
+// intrinsics by default, matching the sm_80 coverage shape.
+
+#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_f16_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<256x256xf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>>
+ %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<256x256xf16>> -> tensor<256x256xf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>> -> tensor<256x256xf16>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf16>, tensor<256x256xf16>) 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_f16_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_16x8x16_F16>
+// CHECK-SAME: promote_operands = [0, 1]
+// CHECK-SAME: reduction = [0, 0, 8]
+// CHECK-SAME: subgroup = [2, 4, 0]
+// CHECK-SAME: workgroup = [64, 64, 0]
+
+// -----
+
+// Test F16 output matmul also selects NV_MMA_SYNC_F16 intrinsic with TileAndFuse.
+
+#pipeline_layout_f16 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_f16_f16() {
+ %cst = arith.constant 0.000000e+00 : f16
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_f16) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_f16) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_f16) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf16>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>> -> tensor<256x256xf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>> -> tensor<256x256xf16>
+ %5 = tensor.empty() : tensor<256x256xf16>
+ %6 = linalg.fill ins(%cst : f16) outs(%5 : tensor<256x256xf16>) -> tensor<256x256xf16>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xf16>, tensor<256x256xf16>) outs(%6 : tensor<256x256xf16>) -> tensor<256x256xf16>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf16> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf16>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_f16_f16(
+// 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_F16_16x8x16_F16>, promote_operands = [0, 1], reduction = [0, 0, 8], subgroup = [2, 4, 0], workgroup = [64, 64, 0]
+
+// -----
+
+// Test that matmul_accumulate sets convert_acc_gemm for NV_MMA_SYNC.
+
+#pipeline_layout_acc = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_accumulate_256x256x256_f16_f32() {
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_acc) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_acc) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_acc) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<readwrite: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<256x256xf16>> -> tensor<256x256xf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xf16>> -> tensor<256x256xf16>
+ %5 = iree_tensor_ext.dispatch.tensor.load %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<256x256xf32>> -> tensor<256x256xf32>
+ %6 = linalg.matmul ins(%3, %4 : tensor<256x256xf16>, tensor<256x256xf16>) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<readwrite:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_accumulate_256x256x256_f16_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: convert_acc_gemm
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_F16>, promote_operands = [0, 1], reduction = [0, 0, 8], subgroup = [2, 4, 0], workgroup = [64, 64, 0]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
index ed215ab..69517d0 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
@@ -62,6 +62,10 @@
// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F16_16x8x16_F16>
+// CHECK-SAME: promote_operands = [0, 1]
+// CHECK-SAME: reduction = [0, 0, 8]
+// CHECK-SAME: subgroup = [2, 4, 0]
+// CHECK-SAME: workgroup = [64, 64, 0]
// -----
@@ -91,6 +95,10 @@
// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
// CHECK-SAME: convert_acc_gemm
// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_F16>
+// CHECK-SAME: promote_operands = [0, 1]
+// CHECK-SAME: reduction = [0, 0, 8]
+// CHECK-SAME: subgroup = [2, 4, 0]
+// CHECK-SAME: workgroup = [64, 64, 0]
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_sm80.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
similarity index 97%
rename from compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_sm80.mlir
rename to compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
index bf861d4..5d99b25 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_sm80.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --iree-gpu-test-target=sm_80 \
+// RUN: iree-opt --split-input-file \
// RUN: --iree-codegen-llvmgpu-nvvm-lowering-pipeline='include-llvm-lowering=false' %s | FileCheck %s
// Test that TileAndFuse pipeline with NV_MMA_SYNC generates nvgpu.mma.sync operations.
@@ -11,7 +11,7 @@
]>
#config = #iree_gpu.lowering_config<{
workgroup = [64, 64, 0],
- reduction = [0, 0, 2],
+ reduction = [0, 0, 8],
subgroup = [2, 4],
mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_F16>,
promote_operands = [0, 1]
@@ -66,7 +66,7 @@
]>
#config_f16 = #iree_gpu.lowering_config<{
workgroup = [64, 64, 0],
- reduction = [0, 0, 2],
+ reduction = [0, 0, 8],
subgroup = [2, 4],
mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F16_16x8x16_F16>,
promote_operands = [0, 1]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_sm80.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_mma_sync.mlir
similarity index 98%
rename from compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_sm80.mlir
rename to compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_mma_sync.mlir
index 543944c..8a80b02 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_sm80.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_vector_distribute_mma_sync.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --iree-gpu-test-target=sm_80 \
+// RUN: iree-opt --split-input-file \
// RUN: --iree-codegen-llvmgpu-nvvm-lowering-pipeline='include-llvm-lowering=false' \
// RUN: %s | FileCheck %s
