[mlir][MemRefToLLVM] Fix crash with unconvertable memory space (#132323)

This PR adds handling when the `memref.alloca` with unconvertable memory
space to prevent a crash. Fixes #131439.

Author: CoTinker

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

@@ -105,8 +105,11 @@ struct AllocaOpLowering : public AllocLikeOpLLVMLowering {
     auto allocaOp = cast<memref::AllocaOp>(op);
     auto elementType =
         typeConverter->convertType(allocaOp.getType().getElementType());
-    unsigned addrSpace =
-        *getTypeConverter()->getMemRefAddressSpace(allocaOp.getType());
+    FailureOr<unsigned> maybeAddressSpace =
+        getTypeConverter()->getMemRefAddressSpace(allocaOp.getType());
+    if (failed(maybeAddressSpace))
+      return std::make_tuple(Value(), Value());
+    unsigned addrSpace = *maybeAddressSpace;
     auto elementPtrType =
         LLVM::LLVMPointerType::get(rewriter.getContext(), addrSpace);
 

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

@@ -654,3 +654,14 @@ func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>,
   memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
   func.return
 }
+
+// -----
+
+// Ensure unconvertable memory space not cause a crash
+
+// CHECK-LABEL: @alloca_unconvertable_memory_space
+func.func @alloca_unconvertable_memory_space() {
+  // CHECK: memref.alloca
+  %alloca = memref.alloca() : memref<1x32x33xi32, #spirv.storage_class<StorageBuffer>>
+  func.return
+}