[Mlir-commits] [mlir] [mlir][nvgpu] Fix crash when handling 0D memref in OptimizeSharedMemoryPass (PR #124517)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/124517

>From 33d92fe505ded06c71fb16d5f837ecda40e515b0 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 27 Jan 2025 15:28:12 +0800
Subject: [PATCH 1/3] [mlir][nvgpu] Fix crash when handling 0D memref in
 OptimizeSharedMemoryPass

This PR adds a check for 0D memref types to prevent a crash.
---
 mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index 31ffacb29256f5..7cef4f60bc011c 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -152,6 +152,9 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   if (!memRefType || !NVGPUDialect::hasSharedMemoryAddressSpace(memRefType))
     return failure();
 
+  if (memRefType.getRank() == 0)
+    return failure();
+
   // Abort if the given value has any sub-views; we do not do any alias
   // analysis.
   bool hasSubView = false;

>From 83175a00ba6fbc11a5d76fa5469add4331a2352f Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 27 Jan 2025 15:29:43 +0800
Subject: [PATCH 2/3] add test

---
 mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
index 5a212815ceb2a1..7477e187286771 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -238,3 +238,13 @@ func.func @abort_if_subview(%arg0: memref<128x128xf16>,
 
   return %mat: vector<1x2xf16>
 }
+
+// -----
+
+// Ensure this case not crash
+
+// CHECK-LABEL: func @test_0_d
+func.func @test_0_d() -> memref<i32, #gpu.address_space<workgroup>> {
+  %alloc = memref.alloc() : memref<i32, #gpu.address_space<workgroup>>
+  return %alloc : memref<i32, #gpu.address_space<workgroup>>
+}

>From 960b5fe0f06b3d75d67a71db40f65b72e1958850 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Mon, 27 Jan 2025 16:17:50 +0800
Subject: [PATCH 3/3] add a comment

---
 mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index 7cef4f60bc011c..72f7296a865f8c 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -152,6 +152,7 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   if (!memRefType || !NVGPUDialect::hasSharedMemoryAddressSpace(memRefType))
     return failure();
 
+  // Not support 0D MemRefs.
   if (memRefType.getRank() == 0)
     return failure();