[Mlir-commits] [mlir] [mlir][DataLayout] Add a default memory space entry to the data layout. (PR #127416)

Fabian Mora llvmlistbot at llvm.org
Tue Mar 11 04:34:40 PDT 2025


https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/127416

>From 48375acf47928ab31d7e6746ea606b7f1f72c3bd Mon Sep 17 00:00:00 2001
From: Fabian Mora <6982088+fabianmcg at users.noreply.github.com>
Date: Sun, 16 Feb 2025 20:54:27 +0000
Subject: [PATCH 1/6] [mlir][DataLayout] Add a default memory space entry to
 the data layout.

This patch adds methods to query the default memory space in the data layout.
This is required as MLIR has no well defined default memory space unlike LLVM which has 0.
While `nullptr` is a candidate for default memory space, the `ptr` dialect
will remove the possibility for `nullptr` memory spaces to avoid undefined
semantics.

This patch also modifies the `DataLayoutTypeInterface::areCompatible` to include
the new DL spec, as it is needed to query the default memory space.
---
 mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td   |  3 ++
 mlir/include/mlir/Dialect/DLTI/DLTIBase.td    |  3 ++
 .../mlir/Interfaces/DataLayoutInterfaces.h    | 10 +++-
 .../mlir/Interfaces/DataLayoutInterfaces.td   | 24 +++++++++-
 mlir/lib/Dialect/DLTI/DLTI.cpp                | 12 ++++-
 mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp      |  6 ++-
 mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp          | 46 +++++++++---------
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp  | 27 +++++++++++
 mlir/test/Dialect/LLVMIR/layout.mlir          |  8 ++++
 mlir/test/Dialect/Ptr/layout.mlir             | 47 ++++++++++++++++---
 .../lib/Dialect/DLTI/TestDataLayoutQuery.cpp  |  5 ++
 mlir/test/lib/Dialect/Test/TestTypes.cpp      |  3 +-
 .../Interfaces/DataLayoutInterfacesTest.cpp   | 19 ++++++++
 13 files changed, 178 insertions(+), 35 deletions(-)

diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
index 007a417653e4c..cf3a72633587d 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
@@ -74,6 +74,9 @@ def DLTI_DataLayoutSpecAttr :
     /// Returns the endiannes identifier.
     StringAttr getEndiannessIdentifier(MLIRContext *context) const;
 
+    /// Returns the default memory space identifier.
+    StringAttr getDefaultMemorySpaceIdentifier(MLIRContext *context) const;
+
     /// Returns the alloca memory space identifier.
     StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
 
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
index 136fe8781e717..1a08bafda54ee 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
@@ -53,6 +53,9 @@ def DLTI_Dialect : Dialect {
     constexpr const static ::llvm::StringLiteral
     kDataLayoutManglingModeKey = "dlti.mangling_mode";
 
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutDefaultMemorySpaceKey = "dlti.default_memory_space";
+
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
     
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 0603bc075ab5e..48af1e5cc8d66 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -74,6 +74,10 @@ getDefaultIndexBitwidth(Type type, const DataLayout &dataLayout,
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultEndianness(DataLayoutEntryInterface entry);
 
+/// Default handler for the default memory space request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultMemorySpace(DataLayoutEntryInterface entry);
+
 /// Default handler for alloca memory space request. Dispatches to the
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
@@ -231,6 +235,9 @@ class DataLayout {
   /// Returns the specified endianness.
   Attribute getEndianness() const;
 
+  /// Returns the default memory space used for memory operations.
+  Attribute getDefaultMemorySpace() const;
+
   /// Returns the memory space used for AllocaOps.
   Attribute getAllocaMemorySpace() const;
 
@@ -285,7 +292,8 @@ class DataLayout {
   mutable std::optional<Attribute> endianness;
   /// Cache for the mangling mode.
   mutable std::optional<Attribute> manglingMode;
-  /// Cache for alloca, global, and program memory spaces.
+  /// Cache for default, alloca, global, and program memory spaces.
+  mutable std::optional<Attribute> defaultMemorySpace;
   mutable std::optional<Attribute> allocaMemorySpace;
   mutable std::optional<Attribute> programMemorySpace;
   mutable std::optional<Attribute> globalMemorySpace;
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 70551c4c5b144..7821a94802914 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -135,6 +135,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
       /*methodName=*/"getEndiannessIdentifier",
       /*args=*/(ins "::mlir::MLIRContext *":$context)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the default memory space identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getDefaultMemorySpaceIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     InterfaceMethod<
       /*description=*/"Returns the alloca memory space identifier.",
       /*retTy=*/"::mlir::StringAttr",
@@ -475,6 +481,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         return ::mlir::detail::getDefaultEndianness(entry);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the memory space used by the ABI computed "
+                      "using the relevant entries. The data layout object "
+                      "can be used for recursive queries.",
+      /*retTy=*/"::mlir::Attribute",
+      /*methodName=*/"getDefaultMemorySpace",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultMemorySpace(entry);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the memory space used by the ABI computed "
                       "using the relevant entries. The data layout object "
@@ -637,11 +655,13 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
     InterfaceMethod<
       /*desc=*/"Returns true if the two lists of entries are compatible, that "
                "is, that `newLayout` spec entries can be nested in an op with "
-               "`oldLayout` spec entries.",
+               "`oldLayout` spec entries. `newSpec` is provided to further"
+               "query data from the spec, e.g. the default address space.",
       /*retTy=*/"bool",
       /*methodName=*/"areCompatible",
       /*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout,
-                    "::mlir::DataLayoutEntryListRef":$newLayout),
+                    "::mlir::DataLayoutEntryListRef":$newLayout,
+                    "::mlir::DataLayoutSpecInterface":$newSpec),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{ return true; }]
     >,
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 92efecc62acd5..16ba44de925d0 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -318,7 +318,8 @@ combineOneSpec(DataLayoutSpecInterface spec,
            "unexpected data layout entry for built-in type");
 
     auto interface = cast<DataLayoutTypeInterface>(typeSample);
-    if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second))
+    if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
+                                 spec))
       return failure();
 
     overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
@@ -379,6 +380,12 @@ DataLayoutSpecAttr::getEndiannessIdentifier(MLIRContext *context) const {
   return Builder(context).getStringAttr(DLTIDialect::kDataLayoutEndiannessKey);
 }
 
+StringAttr DataLayoutSpecAttr::getDefaultMemorySpaceIdentifier(
+    MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutDefaultMemorySpaceKey);
+}
+
 StringAttr
 DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
   return Builder(context).getStringAttr(
@@ -609,7 +616,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
                             << DLTIDialect::kDataLayoutEndiannessBig << "' or '"
                             << DLTIDialect::kDataLayoutEndiannessLittle << "'";
     }
-    if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
+    if (entryName == DLTIDialect::kDataLayoutDefaultMemorySpaceKey ||
+        entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 453b206de294e..aaf30df9482c9 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -350,7 +350,8 @@ LLVMPointerType::getIndexBitwidth(const DataLayout &dataLayout,
 }
 
 bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
-                                    DataLayoutEntryListRef newLayout) const {
+                                    DataLayoutEntryListRef newLayout,
+                                    DataLayoutSpecInterface newSpec) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
@@ -598,7 +599,8 @@ static uint64_t extractStructSpecValue(Attribute attr, StructDLEntryPos pos) {
 }
 
 bool LLVMStructType::areCompatible(DataLayoutEntryListRef oldLayout,
-                                   DataLayoutEntryListRef newLayout) const {
+                                   DataLayoutEntryListRef newLayout,
+                                   DataLayoutSpecInterface newSpec) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp b/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
index 49c8ed977d50c..219a543acc32f 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
@@ -23,13 +23,12 @@ using namespace mlir::ptr;
 
 constexpr const static unsigned kDefaultPointerSizeBits = 64;
 constexpr const static unsigned kBitsInByte = 8;
-constexpr const static unsigned kDefaultPointerAlignment = 8;
-
-static Attribute getDefaultMemorySpace(PtrType ptr) { return nullptr; }
+constexpr const static unsigned kDefaultPointerAlignmentBits = 8;
 
 /// Searches the data layout for the pointer spec, returns nullptr if it is not
 /// found.
-static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type) {
+static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type,
+                               Attribute defaultMemorySpace) {
   for (DataLayoutEntryInterface entry : params) {
     if (!entry.isTypeEntry())
       continue;
@@ -41,20 +40,21 @@ static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type) {
   }
   // If not found, and this is the pointer to the default memory space, assume
   // 64-bit pointers.
-  if (type.getMemorySpace() == getDefaultMemorySpace(type))
+  if (type.getMemorySpace() == defaultMemorySpace)
     return SpecAttr::get(type.getContext(), kDefaultPointerSizeBits,
-                         kDefaultPointerAlignment, kDefaultPointerAlignment,
-                         kDefaultPointerSizeBits);
+                         kDefaultPointerAlignmentBits,
+                         kDefaultPointerAlignmentBits, kDefaultPointerSizeBits);
   return nullptr;
 }
 
 bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
-                            DataLayoutEntryListRef newLayout) const {
+                            DataLayoutEntryListRef newLayout,
+                            DataLayoutSpecInterface newSpec) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
     uint32_t size = kDefaultPointerSizeBits;
-    uint32_t abi = kDefaultPointerAlignment;
+    uint32_t abi = kDefaultPointerAlignmentBits;
     auto newType = llvm::cast<PtrType>(llvm::cast<Type>(newEntry.getKey()));
     const auto *it =
         llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
@@ -65,10 +65,13 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
           return false;
         });
     if (it == oldLayout.end()) {
+      Attribute defaultMemorySpace =
+          mlir::detail::getDefaultMemorySpace(newSpec.getSpecForIdentifier(
+              newSpec.getDefaultMemorySpaceIdentifier(getContext())));
       it = llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
         if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey())) {
           auto ptrTy = llvm::cast<PtrType>(type);
-          return ptrTy.getMemorySpace() == getDefaultMemorySpace(ptrTy);
+          return ptrTy.getMemorySpace() == defaultMemorySpace;
         }
         return false;
       });
@@ -90,43 +93,44 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
 
 uint64_t PtrType::getABIAlignment(const DataLayout &dataLayout,
                                   DataLayoutEntryListRef params) const {
-  if (SpecAttr spec = getPointerSpec(params, *this))
+  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return spec.getAbi() / kBitsInByte;
 
-  return dataLayout.getTypeABIAlignment(
-      get(getContext(), getDefaultMemorySpace(*this)));
+  return dataLayout.getTypeABIAlignment(get(getContext(), defaultMemorySpace));
 }
 
 std::optional<uint64_t>
 PtrType::getIndexBitwidth(const DataLayout &dataLayout,
                           DataLayoutEntryListRef params) const {
-  if (SpecAttr spec = getPointerSpec(params, *this)) {
+  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace)) {
     return spec.getIndex() == SpecAttr::kOptionalSpecValue ? spec.getSize()
                                                            : spec.getIndex();
   }
 
-  return dataLayout.getTypeIndexBitwidth(
-      get(getContext(), getDefaultMemorySpace(*this)));
+  return dataLayout.getTypeIndexBitwidth(get(getContext(), defaultMemorySpace));
 }
 
 llvm::TypeSize PtrType::getTypeSizeInBits(const DataLayout &dataLayout,
                                           DataLayoutEntryListRef params) const {
-  if (SpecAttr spec = getPointerSpec(params, *this))
+  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return llvm::TypeSize::getFixed(spec.getSize());
 
   // For other memory spaces, use the size of the pointer to the default memory
   // space.
-  return dataLayout.getTypeSizeInBits(
-      get(getContext(), getDefaultMemorySpace(*this)));
+  return dataLayout.getTypeSizeInBits(get(getContext(), defaultMemorySpace));
 }
 
 uint64_t PtrType::getPreferredAlignment(const DataLayout &dataLayout,
                                         DataLayoutEntryListRef params) const {
-  if (SpecAttr spec = getPointerSpec(params, *this))
+  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return spec.getPreferred() / kBitsInByte;
 
   return dataLayout.getTypePreferredAlignment(
-      get(getContext(), getDefaultMemorySpace(*this)));
+      get(getContext(), defaultMemorySpace));
 }
 
 LogicalResult PtrType::verifyEntries(DataLayoutEntryListRef entries,
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index 2d098e2db6a96..b7f43ddaae977 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -246,6 +246,17 @@ Attribute mlir::detail::getDefaultEndianness(DataLayoutEntryInterface entry) {
   return entry.getValue();
 }
 
+// Returns the default memory space if specified in the given entry. If the
+// entry is empty the default memory space represented by an empty attribute is
+// returned.
+Attribute mlir::detail::getDefaultMemorySpace(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface()) {
+    return Attribute();
+  }
+
+  return entry.getValue();
+}
+
 // Returns the memory space used for alloca operations if specified in the
 // given entry. If the entry is empty the default memory space represented by
 // an empty attribute is returned.
@@ -605,6 +616,22 @@ mlir::Attribute mlir::DataLayout::getEndianness() const {
   return *endianness;
 }
 
+mlir::Attribute mlir::DataLayout::getDefaultMemorySpace() const {
+  checkValid();
+  if (defaultMemorySpace)
+    return *defaultMemorySpace;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getDefaultMemorySpaceIdentifier(
+            originalLayout.getContext()));
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    defaultMemorySpace = iface.getDefaultMemorySpace(entry);
+  else
+    defaultMemorySpace = detail::getDefaultMemorySpace(entry);
+  return *defaultMemorySpace;
+}
+
 mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
   checkValid();
   if (allocaMemorySpace)
diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index 2dfc289c93065..4e60f991f1bf5 100644
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -6,6 +6,7 @@ module {
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
@@ -18,6 +19,7 @@ module {
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
@@ -30,6 +32,7 @@ module {
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
@@ -50,6 +53,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.ptr<5>, dense<[64, 64, 64]> : vector<3xi64>>,
   #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64, 24]> : vector<4xi64>>,
   #dlti.dl_entry<"dlti.endianness", "little">,
+  #dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>,
   #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
   #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
   #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
@@ -61,6 +65,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alignment = 4
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
+    // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
@@ -73,6 +78,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alignment = 4
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
+    // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
@@ -84,6 +90,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 64
@@ -96,6 +103,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
+    // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 24
diff --git a/mlir/test/Dialect/Ptr/layout.mlir b/mlir/test/Dialect/Ptr/layout.mlir
index 73189a388942a..356c57402651f 100644
--- a/mlir/test/Dialect/Ptr/layout.mlir
+++ b/mlir/test/Dialect/Ptr/layout.mlir
@@ -4,16 +4,18 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 32, abi = 32, preferred = 64>>,
   #dlti.dl_entry<!ptr.ptr<5>,#ptr.spec<size = 64, abi = 64, preferred = 64>>,
   #dlti.dl_entry<!ptr.ptr<4>, #ptr.spec<size = 32, abi = 64, preferred = 64, index = 24>>,
+  #dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>,
   #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
   #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
   #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
   #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>
 >} {
-  // CHECK: @spec
+  // CHECK-LABEL: @spec
   func.func @spec() {
     // CHECK: alignment = 4
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
+    // CHECK: default_memory_space = 7
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
     // CHECK: preferred = 8
@@ -21,19 +23,21 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: size = 4
     // CHECK: stack_alignment = 128
     "test.data_layout_query"() : () -> !ptr.ptr
-    // CHECK: alignment = 4
+    // CHECK: alignment = 1
     // CHECK: alloca_memory_space = 5
-    // CHECK: bitsize = 32
+    // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 7
     // CHECK: global_memory_space = 2
-    // CHECK: index = 32
-    // CHECK: preferred = 8
+    // CHECK: index = 64
+    // CHECK: preferred = 1
     // CHECK: program_memory_space = 3
-    // CHECK: size = 4
+    // CHECK: size = 8
     // CHECK: stack_alignment = 128
     "test.data_layout_query"() : () -> !ptr.ptr<3>
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 64
+    // CHECK: default_memory_space = 7
     // CHECK: global_memory_space = 2
     // CHECK: index = 64
     // CHECK: preferred = 8
@@ -44,6 +48,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: alignment = 8
     // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
+    // CHECK: default_memory_space = 7
     // CHECK: global_memory_space = 2
     // CHECK: index = 24
     // CHECK: preferred = 8
@@ -57,6 +62,36 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // -----
 
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!ptr.ptr<1 : ui64>, #ptr.spec<size = 32, abi = 32, preferred = 32>>,
+  #dlti.dl_entry<"dlti.default_memory_space", 1 : ui64>
+>} {
+  // CHECK-LABEL: @default_memory_space
+  func.func @default_memory_space() {
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: index = 32
+    // CHECK: preferred = 4
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !ptr.ptr
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: index = 32
+    // CHECK: preferred = 4
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !ptr.ptr<1>
+    // CHECK: alignment = 4
+    // CHECK: bitsize = 32
+    // CHECK: index = 32
+    // CHECK: preferred = 4
+    // CHECK: size = 4
+    "test.data_layout_query"() : () -> !ptr.ptr<2>
+    return
+  }
+}
+
+// -----
+
 // expected-error at +2 {{preferred alignment is expected to be at least as large as ABI alignment}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 64, abi = 64, preferred = 32>>
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index eae563b48c119..41a2fe5e4ee75 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -42,6 +42,7 @@ struct TestDataLayoutQuery
       uint64_t preferred = layout.getTypePreferredAlignment(op.getType());
       uint64_t index = layout.getTypeIndexBitwidth(op.getType()).value_or(0);
       Attribute endianness = layout.getEndianness();
+      Attribute defaultMemorySpace = layout.getDefaultMemorySpace();
       Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
       Attribute manglingMode = layout.getManglingMode();
       Attribute programMemorySpace = layout.getProgramMemorySpace();
@@ -69,6 +70,10 @@ struct TestDataLayoutQuery
            builder.getNamedAttr("endianness", endianness == Attribute()
                                                   ? builder.getStringAttr("")
                                                   : endianness),
+           builder.getNamedAttr("default_memory_space",
+                                defaultMemorySpace == Attribute()
+                                    ? builder.getUI32IntegerAttr(0)
+                                    : defaultMemorySpace),
            builder.getNamedAttr("alloca_memory_space",
                                 allocaMemorySpace == Attribute()
                                     ? builder.getUI32IntegerAttr(0)
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index 0c237440834ef..e0d8cdda86159 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -284,7 +284,8 @@ TestTypeWithLayoutType::getIndexBitwidth(const DataLayout &dataLayout,
 }
 
 bool TestTypeWithLayoutType::areCompatible(
-    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout) const {
+    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
+    DataLayoutSpecInterface newSpec) const {
   unsigned old = extractKind(oldLayout, "alignment");
   return old == 1 || extractKind(newLayout, "alignment") <= old;
 }
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index 3d87948c5b5af..fd81f3021aa9b 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -23,6 +23,8 @@ using namespace mlir;
 namespace {
 constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
 constexpr static llvm::StringLiteral kEndiannesKeyName = "dltest.endianness";
+constexpr static llvm::StringLiteral kDefaultKeyName =
+    "dltest.default_memory_space";
 constexpr static llvm::StringLiteral kAllocaKeyName =
     "dltest.alloca_memory_space";
 constexpr static llvm::StringLiteral kManglingModeKeyName =
@@ -82,6 +84,9 @@ struct CustomDataLayoutSpec
   StringAttr getEndiannessIdentifier(MLIRContext *context) const {
     return Builder(context).getStringAttr(kEndiannesKeyName);
   }
+  StringAttr getDefaultMemorySpaceIdentifier(MLIRContext *context) const {
+    return Builder(context).getStringAttr(kDefaultKeyName);
+  }
   StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
     return Builder(context).getStringAttr(kAllocaKeyName);
   }
@@ -205,6 +210,15 @@ struct SingleQueryType
     return Attribute();
   }
 
+  Attribute getDefaultMemorySpace(DataLayoutEntryInterface entry) {
+    static bool executed = false;
+    if (executed)
+      llvm::report_fatal_error("repeated call");
+
+    executed = true;
+    return Attribute();
+  }
+
   Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) {
     static bool executed = false;
     if (executed)
@@ -475,6 +489,7 @@ module {}
   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
 
   EXPECT_EQ(layout.getEndianness(), Attribute());
+  EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
   EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
@@ -508,6 +523,7 @@ TEST(DataLayout, NullSpec) {
   EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u);
 
   EXPECT_EQ(layout.getEndianness(), Attribute());
+  EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
   EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
@@ -549,6 +565,7 @@ TEST(DataLayout, EmptySpec) {
   EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u);
 
   EXPECT_EQ(layout.getEndianness(), Attribute());
+  EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
   EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
@@ -572,6 +589,7 @@ TEST(DataLayout, SpecWithEntries) {
   #dlti.dl_entry<i16, 6>,
   #dlti.dl_entry<index, 42>,
   #dlti.dl_entry<"dltest.endianness", "little">,
+  #dlti.dl_entry<"dltest.default_memory_space", 1 : i32>,
   #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>,
   #dlti.dl_entry<"dltest.program_memory_space", 3 : i32>,
   #dlti.dl_entry<"dltest.global_memory_space", 2 : i32>,
@@ -609,6 +627,7 @@ TEST(DataLayout, SpecWithEntries) {
   EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
 
   EXPECT_EQ(layout.getEndianness(), Builder(&ctx).getStringAttr("little"));
+  EXPECT_EQ(layout.getDefaultMemorySpace(), Builder(&ctx).getI32IntegerAttr(1));
   EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5));
   EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3));
   EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2));

>From 7ffa3bc180a56e4e53564c3f5911ee99f6a3de1b Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Mon, 17 Feb 2025 09:29:01 -0500
Subject: [PATCH 2/6] address reviewer comments.

---
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index b7f43ddaae977..7afa6ac3a00e4 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -250,9 +250,8 @@ Attribute mlir::detail::getDefaultEndianness(DataLayoutEntryInterface entry) {
 // entry is empty the default memory space represented by an empty attribute is
 // returned.
 Attribute mlir::detail::getDefaultMemorySpace(DataLayoutEntryInterface entry) {
-  if (entry == DataLayoutEntryInterface()) {
+  if (!entry)
     return Attribute();
-  }
 
   return entry.getValue();
 }

>From e8c8e66fa3c217503503e305e226ca071a50f2b1 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Thu, 27 Feb 2025 10:28:42 -0500
Subject: [PATCH 3/6] Update
 mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

Co-authored-by: Christian Ulmann <christianulmann at gmail.com>
---
 mlir/include/mlir/Interfaces/DataLayoutInterfaces.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 7821a94802914..ba67dbe640b8a 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -656,7 +656,7 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
       /*desc=*/"Returns true if the two lists of entries are compatible, that "
                "is, that `newLayout` spec entries can be nested in an op with "
                "`oldLayout` spec entries. `newSpec` is provided to further"
-               "query data from the spec, e.g. the default address space.",
+               "query data from the spec, e.g., the default address space.",
       /*retTy=*/"bool",
       /*methodName=*/"areCompatible",
       /*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout,

>From a3a55d6f03f9460e4e8ccd5a965a81839348c019 Mon Sep 17 00:00:00 2001
From: Fabian Mora <6982088+fabianmcg at users.noreply.github.com>
Date: Sun, 2 Mar 2025 15:06:21 +0000
Subject: [PATCH 4/6] address reviewer comments

---
 .../mlir/Interfaces/DataLayoutInterfaces.h    |  2 +
 .../mlir/Interfaces/DataLayoutInterfaces.td   |  8 ++--
 mlir/lib/Dialect/DLTI/DLTI.cpp                | 44 ++++++++++---------
 mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp      | 14 +++---
 mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp          |  8 ++--
 mlir/test/lib/Dialect/Test/TestTypes.cpp      |  3 +-
 6 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 48af1e5cc8d66..50d85a8ccb1be 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -34,6 +34,8 @@ using DataLayoutEntryList = llvm::SmallVector<DataLayoutEntryInterface, 4>;
 using DataLayoutEntryListRef = llvm::ArrayRef<DataLayoutEntryInterface>;
 using TargetDeviceSpecListRef = llvm::ArrayRef<TargetDeviceSpecInterface>;
 using TargetDeviceSpecEntry = std::pair<StringAttr, TargetDeviceSpecInterface>;
+using DataLayoutIdentifiedEntryMap =
+    ::llvm::DenseMap<::mlir::StringAttr, ::mlir::DataLayoutEntryInterface>;
 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
 class ModuleOp;
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index ba67dbe640b8a..a02258dbad4eb 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -655,13 +655,15 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
     InterfaceMethod<
       /*desc=*/"Returns true if the two lists of entries are compatible, that "
                "is, that `newLayout` spec entries can be nested in an op with "
-               "`oldLayout` spec entries. `newSpec` is provided to further"
-               "query data from the spec, e.g., the default address space.",
+               "`oldLayout` spec entries. `newSpec` and `identified` are"
+               "provided to further query data from the combined spec, e.g.,"
+               "the default address space.",
       /*retTy=*/"bool",
       /*methodName=*/"areCompatible",
       /*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout,
                     "::mlir::DataLayoutEntryListRef":$newLayout,
-                    "::mlir::DataLayoutSpecInterface":$newSpec),
+                    "::mlir::DataLayoutSpecInterface":$newSpec,
+                    "const ::mlir::DataLayoutIdentifiedEntryMap&":$identified),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{ return true; }]
     >,
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 16ba44de925d0..0cfad2ba3a522 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -305,26 +305,8 @@ combineOneSpec(DataLayoutSpecInterface spec,
   DenseMap<StringAttr, DataLayoutEntryInterface> newEntriesForID;
   spec.bucketEntriesByType(newEntriesForType, newEntriesForID);
 
-  // Try overwriting the old entries with the new ones.
-  for (auto &kvp : newEntriesForType) {
-    if (!entriesForType.count(kvp.first)) {
-      entriesForType[kvp.first] = std::move(kvp.second);
-      continue;
-    }
-
-    Type typeSample = cast<Type>(kvp.second.front().getKey());
-    assert(&typeSample.getDialect() !=
-               typeSample.getContext()->getLoadedDialect<BuiltinDialect>() &&
-           "unexpected data layout entry for built-in type");
-
-    auto interface = cast<DataLayoutTypeInterface>(typeSample);
-    if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
-                                 spec))
-      return failure();
-
-    overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
-  }
-
+  // Combine non-Type DL entries first so they are visible to the
+  // `type.areCompatible` method, allowing to query global properties.
   for (const auto &kvp : newEntriesForID) {
     StringAttr id = cast<StringAttr>(kvp.second.getKey());
     Dialect *dialect = id.getReferencedDialect();
@@ -333,7 +315,7 @@ combineOneSpec(DataLayoutSpecInterface spec,
       continue;
     }
 
-    // Attempt to combine the enties using the dialect interface. If the
+    // Attempt to combine the entries using the dialect interface. If the
     // dialect is not loaded for some reason, use the default combinator
     // that conservatively accepts identical entries only.
     entriesForID[id] =
@@ -345,6 +327,26 @@ combineOneSpec(DataLayoutSpecInterface spec,
       return failure();
   }
 
+  // Try overwriting the old entries with the new ones.
+  for (auto &kvp : newEntriesForType) {
+    if (!entriesForType.count(kvp.first)) {
+      entriesForType[kvp.first] = std::move(kvp.second);
+      continue;
+    }
+
+    Type typeSample = cast<Type>(kvp.second.front().getKey());
+    assert(&typeSample.getDialect() !=
+               typeSample.getContext()->getLoadedDialect<BuiltinDialect>() &&
+           "unexpected data layout entry for built-in type");
+
+    auto interface = cast<DataLayoutTypeInterface>(typeSample);
+    if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
+                                 spec, entriesForID))
+      return failure();
+
+    overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
+  }
+
   return success();
 }
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index aaf30df9482c9..8f39ede721c92 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
@@ -349,9 +349,10 @@ LLVMPointerType::getIndexBitwidth(const DataLayout &dataLayout,
   return dataLayout.getTypeIndexBitwidth(get(getContext()));
 }
 
-bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
-                                    DataLayoutEntryListRef newLayout,
-                                    DataLayoutSpecInterface newSpec) const {
+bool LLVMPointerType::areCompatible(
+    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
+    DataLayoutSpecInterface newSpec,
+    const DataLayoutIdentifiedEntryMap &map) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
@@ -598,9 +599,10 @@ static uint64_t extractStructSpecValue(Attribute attr, StructDLEntryPos pos) {
       .getValues<uint64_t>()[static_cast<size_t>(pos)];
 }
 
-bool LLVMStructType::areCompatible(DataLayoutEntryListRef oldLayout,
-                                   DataLayoutEntryListRef newLayout,
-                                   DataLayoutSpecInterface newSpec) const {
+bool LLVMStructType::areCompatible(
+    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
+    DataLayoutSpecInterface newSpec,
+    const DataLayoutIdentifiedEntryMap &map) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp b/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
index 219a543acc32f..1fc424c05f6c7 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp
@@ -49,7 +49,8 @@ static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type,
 
 bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
                             DataLayoutEntryListRef newLayout,
-                            DataLayoutSpecInterface newSpec) const {
+                            DataLayoutSpecInterface newSpec,
+                            const DataLayoutIdentifiedEntryMap &map) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
@@ -65,9 +66,8 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
           return false;
         });
     if (it == oldLayout.end()) {
-      Attribute defaultMemorySpace =
-          mlir::detail::getDefaultMemorySpace(newSpec.getSpecForIdentifier(
-              newSpec.getDefaultMemorySpaceIdentifier(getContext())));
+      Attribute defaultMemorySpace = mlir::detail::getDefaultMemorySpace(
+          map.lookup(newSpec.getDefaultMemorySpaceIdentifier(getContext())));
       it = llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
         if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey())) {
           auto ptrTy = llvm::cast<PtrType>(type);
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index e0d8cdda86159..5c784dcee6e15 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -285,7 +285,8 @@ TestTypeWithLayoutType::getIndexBitwidth(const DataLayout &dataLayout,
 
 bool TestTypeWithLayoutType::areCompatible(
     DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
-    DataLayoutSpecInterface newSpec) const {
+    DataLayoutSpecInterface newSpec,
+    const DataLayoutIdentifiedEntryMap &map) const {
   unsigned old = extractKind(oldLayout, "alignment");
   return old == 1 || extractKind(newLayout, "alignment") <= old;
 }

>From c3224ad204cc81229046b8c53496bf8e37b6a6fe Mon Sep 17 00:00:00 2001
From: fabian <6982088+fabianmcg at users.noreply.github.com>
Date: Fri, 7 Mar 2025 10:39:54 -0800
Subject: [PATCH 5/6] add TODO

---
 mlir/lib/Dialect/DLTI/DLTI.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 0cfad2ba3a522..c81da435618dd 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -340,6 +340,7 @@ combineOneSpec(DataLayoutSpecInterface spec,
            "unexpected data layout entry for built-in type");
 
     auto interface = cast<DataLayoutTypeInterface>(typeSample);
+    // TODO: Revisit this method and call once issue #130321 gets resolved.
     if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
                                  spec, entriesForID))
       return failure();

>From 6d636ff1566c4b0e48315740bd18994a6c3ad0a4 Mon Sep 17 00:00:00 2001
From: fabian <6982088+fabianmcg at users.noreply.github.com>
Date: Tue, 11 Mar 2025 04:34:26 -0700
Subject: [PATCH 6/6] address reviewer comments

---
 mlir/include/mlir/Interfaces/DataLayoutInterfaces.td | 3 ++-
 mlir/lib/Dialect/DLTI/DLTI.cpp                       | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index a02258dbad4eb..818b441b9a770 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -657,7 +657,8 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
                "is, that `newLayout` spec entries can be nested in an op with "
                "`oldLayout` spec entries. `newSpec` and `identified` are"
                "provided to further query data from the combined spec, e.g.,"
-               "the default address space.",
+               "the default address space. TODO: Revisit this method once"
+               "https://github.com/llvm/llvm-project/issues/130321 gets solved",
       /*retTy=*/"bool",
       /*methodName=*/"areCompatible",
       /*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout,
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index c81da435618dd..206d9f43a44a4 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -340,7 +340,8 @@ combineOneSpec(DataLayoutSpecInterface spec,
            "unexpected data layout entry for built-in type");
 
     auto interface = cast<DataLayoutTypeInterface>(typeSample);
-    // TODO: Revisit this method and call once issue #130321 gets resolved.
+    // TODO: Revisit this method and call once
+    // https://github.com/llvm/llvm-project/issues/130321 gets resolved.
     if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
                                  spec, entriesForID))
       return failure();



More information about the Mlir-commits mailing list