[mlir][nvgpu] Fix crash when handling 0D memref in OptimizeSharedMemoryPass (llvm#124517)

This PR adds a check for 0D memref types to prevent a crash. Fixes
llvm#119855.

Author: CoTinker

mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp

@@ -152,6 +152,10 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   if (!memRefType || !NVGPUDialect::hasSharedMemoryAddressSpace(memRefType))
     return failure();
 
+  // Not support 0D MemRefs.
+  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;

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>>
+}