[mlir][memref] Check memory space before lowering alloc ops (#134427)

Check the memory space before lowering allocation ops, instead of
starting the lowering and then rolling back the pattern when the memory
space was found to be incompatible with LLVM.

Note: This is in preparation of the One-Shot Dialect Conversion
refactoring.

Note: `isConvertibleAndHasIdentityMaps` now also checks the memory
space.

Author: matthias-springer

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

@@ -75,8 +75,8 @@ class ConvertToLLVMPattern : public ConversionPattern {
                              ValueRange indices,
                              ConversionPatternRewriter &rewriter) const;
 
-  /// Returns if the given memref has identity maps and the element type is
-  /// convertible to LLVM.
+  /// Returns if the given memref type is convertible to LLVM and has an
+  /// identity layout map.
   bool isConvertibleAndHasIdentityMaps(MemRefType type) const;
 
   /// Returns the type of a pointer to an element of the memref.

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -99,9 +99,9 @@ Value ConvertToLLVMPattern::getStridedElementPtr(
 // only support memrefs with identity maps.
 bool ConvertToLLVMPattern::isConvertibleAndHasIdentityMaps(
     MemRefType type) const {
-  if (!typeConverter->convertType(type.getElementType()))
+  if (!type.getLayout().isIdentity())
     return false;
-  return type.getLayout().isIdentity();
+  return static_cast<bool>(typeConverter->convertType(type));
 }
 
 Type ConvertToLLVMPattern::getElementPtrType(MemRefType type) const {

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

@@ -73,12 +73,7 @@ std::tuple<Value, Value> AllocationOpLLVMLowering::allocateBufferManuallyAlign(
   MemRefType memRefType = getMemRefResultType(op);
   // Allocate the underlying buffer.
   Type elementPtrType = this->getElementPtrType(memRefType);
-  if (!elementPtrType) {
-    emitError(loc, "conversion of memref memory space ")
-        << memRefType.getMemorySpace()
-        << " to integer address space "
-           "failed. Consider adding memory space conversions.";
-  }
+  assert(elementPtrType && "could not compute element ptr type");
   FailureOr<LLVM::LLVMFuncOp> allocFuncOp = getNotalignedAllocFn(
       getTypeConverter(), op->getParentWithTrait<OpTrait::SymbolTable>(),
       getIndexType());

mlir/test/Conversion/MemRefToLLVM/invalid.mlir

@@ -22,7 +22,7 @@ func.func @bad_address_space(%a: memref<2xindex, "foo">) {
 
 // CHECK-LABEL: @invalid_int_conversion
 func.func @invalid_int_conversion() {
-     // expected-error@+1 {{conversion of memref memory space 1 : ui64 to integer address space failed. Consider adding memory space conversions.}}
+     // expected-error@unknown{{conversion of memref memory space 1 : ui64 to integer address space failed. Consider adding memory space conversions.}}
      %alloc = memref.alloc() {alignment = 64 : i64} : memref<10xf32, 1 : ui64> 
     return
 }
@@ -32,7 +32,6 @@ func.func @invalid_int_conversion() {
 // expected-error@unknown{{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() {
-  // expected-error@+1{{conversion of memref memory space #gpu.address_space<workgroup> to integer address space failed. Consider adding memory space conversions}}
   %alloc = memref.alloc() : memref<1x32x33xi32, #gpu.address_space<workgroup>>
   %alloc1 = memref.alloc() : memref<i32>
   %c0 = arith.constant 0 : index