[llvm-branch-commits] [mlir] [mlir][memref] Check memory space before lowering alloc ops (PR #134427)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/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.

Depends on #134421.


>From bd104624a51dc315b94f651271b95b8b438a8146 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Fri, 4 Apr 2025 19:59:28 +0200
Subject: [PATCH] [mlir][memref] Check memory space before lowering alloc ops

---
 mlir/include/mlir/Conversion/LLVMCommon/Pattern.h        | 8 +++++---
 mlir/lib/Conversion/LLVMCommon/Pattern.cpp               | 9 ++++++---
 mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp | 7 +------
 mlir/test/Conversion/MemRefToLLVM/invalid.mlir           | 3 +--
 4 files changed, 13 insertions(+), 14 deletions(-)

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 at +1 {{conversion of memref memory space 1 : ui64 to integer address space failed. Consider adding memory space conversions.}}
+     // expected-error at 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 at 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 at +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