[mlir][memref] Check memory space before lowering alloc ops

Author: matthias-springer

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

@@ -75,9 +75,11 @@ 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.
-  bool isConvertibleAndHasIdentityMaps(MemRefType type) const;
+  /// Returns if the given memref type is convertible to LLVM and has an
+  /// identity layout map. If `verifyMemorySpace` is set to "false", the memory
+  /// space of the memref type is ignored.
+  bool isConvertibleAndHasIdentityMaps(MemRefType type,
+                                       bool verifyMemorySpace = true) const;
 
   /// Returns the type of a pointer to an element of the memref.
   Type getElementPtrType(MemRefType type) const;

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -98,10 +98,13 @@ Value ConvertToLLVMPattern::getStridedElementPtr(
 // Check if the MemRefType `type` is supported by the lowering. We currently
 // only support memrefs with identity maps.
 bool ConvertToLLVMPattern::isConvertibleAndHasIdentityMaps(
-    MemRefType type) const {
-  if (!typeConverter->convertType(type.getElementType()))
+    MemRefType type, bool verifyMemorySpace) const {
+  if (!type.getLayout().isIdentity())
     return false;
-  return type.getLayout().isIdentity();
+  // If the memory space should not be verified, just check the element type.
+  Type typeToVerify =
+      verifyMemorySpace ? static_cast<Type>(type) : type.getElementType();
+  return static_cast<bool>(typeConverter->convertType(typeToVerify));
 }
 
 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