[mlir][MemRefToLLVM] Fix crash with unconvertable memory space

[CoTinker]

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

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

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

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+5-2)
• (added) mlir/test/Conversion/MemRefToLLVM/issue-131439.mlir (+15)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index fe0ee11d84adb..cb4317ef1bcec 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/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);
 
diff --git a/mlir/test/Conversion/MemRefToLLVM/issue-131439.mlir b/mlir/test/Conversion/MemRefToLLVM/issue-131439.mlir
new file mode 100644
index 0000000000000..d5aa75fb45f52
--- /dev/null
+++ b/mlir/test/Conversion/MemRefToLLVM/issue-131439.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-opt %s -finalize-memref-to-llvm 2>&1 | FileCheck %s
+// Since the error is at an unknown location, we use FileCheck instead of
+// -verify-diagnostics here
+
+// CHECK: conversion of memref memory space #spirv.storage_class<StorageBuffer> to integer address space failed. Consider adding memory space conversions.
+// CHECK-LABEL: @issue_131439
+func.func @issue_131439(%arg0 : memref<i32>) {
+  // CHECK: memref.alloca
+  %alloca = memref.alloca() : memref<1x32x33xi32, #spirv.storage_class<StorageBuffer>>
+  %c0 = arith.constant 0 : index
+  // CHECK: memref.load
+  %0 = memref.load %alloca[%c0, %c0, %c0] : memref<1x32x33xi32, #spirv.storage_class<StorageBuffer>>
+  memref.store %0, %arg0[] : memref<i32>
+  func.return
+}

[krzysz00]

Approved, thanks

[CoTinker]

Thanks for your review.