[Mlir-commits] [mlir] [mlir][MemRefToLLVM] Fix crashes with unconvertable memory spaces (PR #70694)

[Krzysztof Drewniak]

https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/70694

>From e622a0730160294d1624332696192139f71bc9b4 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 30 Oct 2023 17:32:08 +0000
Subject: [PATCH] [mlir][MemRefToLLVM] Fix crashes with unconvertable memory
 spaces

Fixes #70160

The issue is resolved by:
1. Changing the call to address space conversion to use the correct
return type, preventing the code from moving past the if and into the
crashing optional dereference.
2. Adding handling to the AllocLikeOp rewriter for the case where the
underlying buffer allocation fails.
---
 .../MemRefToLLVM/AllocLikeConversion.cpp          | 13 +++++++++++--
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp |  6 ++++--
 mlir/test/Conversion/MemRefToLLVM/invalid.mlir    |  4 +---
 .../test/Conversion/MemRefToLLVM/issue-70160.mlir | 15 +++++++++++++++
 4 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 mlir/test/Conversion/MemRefToLLVM/issue-70160.mlir

diff --git a/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp b/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
index a2a426e3c293175..6286b32b3704279 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
@@ -59,7 +59,11 @@ static Value castAllocFuncResult(ConversionPatternRewriter &rewriter,
                                  MemRefType memRefType, Type elementPtrType,
                                  const LLVMTypeConverter &typeConverter) {
   auto allocatedPtrTy = cast<LLVM::LLVMPointerType>(allocatedPtr.getType());
-  unsigned memrefAddrSpace = *typeConverter.getMemRefAddressSpace(memRefType);
+  FailureOr<unsigned> maybeMemrefAddrSpace =
+      typeConverter.getMemRefAddressSpace(memRefType);
+  if (failed(maybeMemrefAddrSpace))
+    return Value();
+  unsigned memrefAddrSpace = *maybeMemrefAddrSpace;
   if (allocatedPtrTy.getAddressSpace() != memrefAddrSpace)
     allocatedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
         loc,
@@ -91,7 +95,8 @@ std::tuple<Value, Value> AllocationOpLLVMLowering::allocateBufferManuallyAlign(
   Value allocatedPtr =
       castAllocFuncResult(rewriter, loc, results.getResult(), memRefType,
                           elementPtrType, *getTypeConverter());
-
+  if (!allocatedPtr)
+    return std::make_tuple(Value(), Value());
   Value alignedPtr = allocatedPtr;
   if (alignment) {
     // Compute the aligned pointer.
@@ -183,6 +188,10 @@ LogicalResult AllocLikeOpLLVMLowering::matchAndRewrite(
   auto [allocatedPtr, alignedPtr] =
       this->allocateBuffer(rewriter, loc, size, op);
 
+  if (!allocatedPtr || !alignedPtr)
+    return rewriter.notifyMatchFailure(loc,
+                                       "underlying buffer allocation failed");
+
   // Create the MemRef descriptor.
   auto memRefDescriptor = this->createMemRefDescriptor(
       loc, memRefType, allocatedPtr, alignedPtr, sizes, strides, rewriter);
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 61bd23f12601c79..e9efcedec0e14e4 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -76,6 +76,8 @@ struct AlignedAllocOpLowering : public AllocLikeOpLLVMLowering {
         rewriter, loc, sizeBytes, op, &defaultLayout,
         alignedAllocationGetAlignment(rewriter, loc, cast<memref::AllocOp>(op),
                                       &defaultLayout));
+    if (!ptr)
+      return std::make_tuple(Value(), Value());
     return std::make_tuple(ptr, ptr);
   }
 
@@ -560,9 +562,9 @@ struct GetGlobalMemrefOpLowering : public AllocLikeOpLLVMLowering {
 
     // This is called after a type conversion, which would have failed if this
     // call fails.
-    std::optional<unsigned> maybeAddressSpace =
+    FailureOr<unsigned> maybeAddressSpace =
         getTypeConverter()->getMemRefAddressSpace(type);
-    if (!maybeAddressSpace)
+    if (failed(maybeAddressSpace))
       return std::make_tuple(Value(), Value());
     unsigned memSpace = *maybeAddressSpace;
 
diff --git a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir b/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
index 786e6f256219934..40dd75af1dd7703 100644
--- a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -finalize-memref-to-llvm -split-input-file 2>&1 | FileCheck %s
+// 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
 // -veri-y-diagnostics here
 
@@ -10,5 +10,3 @@ func.func @bad_address_space(%a: memref<2xindex, "foo">) {
     memref.store %c0, %a[%c0] : memref<2xindex, "foo">
     return
 }
-
-// -----
diff --git a/mlir/test/Conversion/MemRefToLLVM/issue-70160.mlir b/mlir/test/Conversion/MemRefToLLVM/issue-70160.mlir
new file mode 100644
index 000000000000000..4f94e6375131b1b
--- /dev/null
+++ b/mlir/test/Conversion/MemRefToLLVM/issue-70160.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
+// -veri-y-diagnostics here
+
+// CHECK: conversion of memref memory space #gpu.address_space<workgroup> to integer address space failed. Consider adding memory space conversions
+// CHECK-LABEL: @issue_70160
+func.func @issue_70160() {
+  %alloc = memref.alloc() : memref<1x32x33xi32, #gpu.address_space<workgroup>>
+  %alloc1 = memref.alloc() : memref<i32>
+  %c0 = arith.constant 0 : index
+  // CHECK: memref.load
+  %0 = memref.load %alloc[%c0, %c0, %c0] : memref<1x32x33xi32, #gpu.address_space<workgroup>>
+  memref.store %0, %alloc1[] : memref<i32>
+  func.return
+}