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

[llvmlistbot at llvm.org]

Author: Krzysztof Drewniak
Date: 2023-10-31T09:01:43-05:00
New Revision: 73c6248cc2cc3acd01c3580bfdc64825c09a0fd6

URL: https://github.com/llvm/llvm-project/commit/73c6248cc2cc3acd01c3580bfdc64825c09a0fd6
DIFF: https://github.com/llvm/llvm-project/commit/73c6248cc2cc3acd01c3580bfdc64825c09a0fd6.diff

LOG: [mlir][MemRefToLLVM] Fix crashes with unconvertable memory spaces (#70694)

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.

Added: 
    mlir/test/Conversion/MemRefToLLVM/issue-70160.mlir

Modified: 
    mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
    mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
    mlir/test/Conversion/MemRefToLLVM/invalid.mlir

Removed: 
    


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