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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-nvgpu

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This PR adds a check for 0D memref types to prevent a crash. Fixes #<!-- -->119855.

---
Full diff: https://github.com/llvm/llvm-project/pull/124517.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp (+3) 
- (modified) mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir (+10) 


``````````diff
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;
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>>
+}

``````````

</details>


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