[Mlir-commits] [mlir] [mlir][nvgpu] Add roundtrip tests for warpgroup MMA operations (PR #199272)

[llvmlistbot at llvm.org]

https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/199272

>From 37c941385c577b2fd4341aab226531510f7c8dfe Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Fri, 22 May 2026 20:38:16 +0200
Subject: [PATCH 1/2] [mlir][nvgpu] Add roundtrip tests for warpgroup MMA
 operations

Add missing roundtrip tests for nvgpu.warpgroup.mma.init.accumulator
and nvgpu.warpgroup.mma.store operations in the NVGPU dialect.
The tests verify that operand and result types survive the
parser/printer roundtrip.
---
 mlir/test/Dialect/NVGPU/roundtrip.mlir | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/test/Dialect/NVGPU/roundtrip.mlir b/mlir/test/Dialect/NVGPU/roundtrip.mlir
index ad516b4d2c200..51f6340041513 100644
--- a/mlir/test/Dialect/NVGPU/roundtrip.mlir
+++ b/mlir/test/Dialect/NVGPU/roundtrip.mlir
@@ -69,3 +69,23 @@ func.func @async_cp(%dst : memref<2x7x5xf32, 3>, %src : memref<4x5xf32>){
   nvgpu.device_async_wait %token {numGroups = 1 : i32}
   return
 }
+
+// CHECK-LABEL: func @warpgroup_mma_init_accumulator
+func.func @warpgroup_mma_init_accumulator() {
+  // CHECK: %{{.*}} = nvgpu.warpgroup.mma.init.accumulator
+  // CHECK-SAME: -> <fragmented = vector<64x128xf32>>
+  %acc = nvgpu.warpgroup.mma.init.accumulator
+      -> !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>
+  return
+}
+
+// CHECK-LABEL: func @warpgroup_mma_store
+func.func @warpgroup_mma_store(
+    %acc: !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>,
+    %dst: memref<64x128xf32, 3>) {
+  // CHECK: nvgpu.warpgroup.mma.store %{{.*}}, %{{.*}} : <fragmented = vector<64x128xf32>> to memref<64x128xf32, 3>
+  nvgpu.warpgroup.mma.store %acc, %dst :
+      !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>
+      to memref<64x128xf32, 3>
+  return
+}
\ No newline at end of file

>From b41984af06044cdc9c7ff941bbc217031c75aed2 Mon Sep 17 00:00:00 2001
From: Vaisman <vasili.svirski at gmail.com>
Date: Wed, 27 May 2026 10:07:01 +0200
Subject: [PATCH 2/2] Fix missing newline at end of roundtrip.mlir

---
 mlir/test/Dialect/NVGPU/roundtrip.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/NVGPU/roundtrip.mlir b/mlir/test/Dialect/NVGPU/roundtrip.mlir
index 51f6340041513..77586fc9b7abf 100644
--- a/mlir/test/Dialect/NVGPU/roundtrip.mlir
+++ b/mlir/test/Dialect/NVGPU/roundtrip.mlir
@@ -88,4 +88,4 @@ func.func @warpgroup_mma_store(
       !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>
       to memref<64x128xf32, 3>
   return
-}
\ No newline at end of file
+}