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

[matthias-springer]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

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.

Depends on #134421.

---

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

4 Files Affected:

• (modified) mlir/include/mlir/Conversion/LLVMCommon/Pattern.h (+5-3)
• (modified) mlir/lib/Conversion/LLVMCommon/Pattern.cpp (+6-3)
• (modified) mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp (+1-6)
• (modified) mlir/test/Conversion/MemRefToLLVM/invalid.mlir (+1-2)

diff --git a/mlir/include/mlir/Conversion/LLVMCommon/Pattern.h b/mlir/include/mlir/Conversion/LLVMCommon/Pattern.h
index c65f7d7217be5..6f7811acec939 100644
--- a/mlir/include/mlir/Conversion/LLVMCommon/Pattern.h
+++ b/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;
diff --git a/mlir/lib/Conversion/LLVMCommon/Pattern.cpp b/mlir/lib/Conversion/LLVMCommon/Pattern.cpp
index 71b68619cc793..d11de1f44250c 100644
--- a/mlir/lib/Conversion/LLVMCommon/Pattern.cpp
+++ b/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 {
diff --git a/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp b/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
index c5b2e83df93dc..bad209a4ddecf 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp
+++ b/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());
diff --git a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir b/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
index 61c67005a08fc..0d04bba96bcdb 100644
--- a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
+++ b/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

[zero9178]

(Non actionable)
This is a bit unfortunate regarding the quality of the error reporting, but is more of a pre-existing issue with the conversion functions (and the type converter infrastructure). I like the fact all type conversion related error handling is actually in the type conversions now.

[matthias-springer]

I'm thinking of maybe even removing this error entirely (not in this PR). This error is produced in the type converter. We don't usually produce errors there.

There are other ways to signal this to the user: when the type converter fails, the pattern does not apply. Either because a hand-written converter.convertType fails or because the values of the adaptor could not be computed. In either case, -debug would show that (notifyMatchFailure).

[zero9178]

That sounds great to me, I'd support that