[mlir][DataLayout] Add a default memory space entry to the data layout.

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;
 

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";
 

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;
@@ -74,6 +76,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 +237,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 +294,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;

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,16 @@ 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` and `identified` are"
+               "provided to further query data from the combined spec, e.g.,"
+               "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,
-                    "::mlir::DataLayoutEntryListRef":$newLayout),
+                    "::mlir::DataLayoutEntryListRef":$newLayout,
+                    "::mlir::DataLayoutSpecInterface":$newSpec,
+                    "const ::mlir::DataLayoutIdentifiedEntryMap&":$identified),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{ return true; }]
     >,

mlir/lib/Dialect/DLTI/DLTI.cpp

@@ -305,25 +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))
-      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();
@@ -332,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] =
@@ -344,6 +327,28 @@ 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);
+    // 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();
+
+    overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
+  }
+
   return success();
 }
 
@@ -379,6 +384,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 +620,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 ||

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

@@ -349,8 +349,10 @@ LLVMPointerType::getIndexBitwidth(const DataLayout &dataLayout,
   return dataLayout.getTypeIndexBitwidth(get(getContext()));
 }
 
-bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
-                                    DataLayoutEntryListRef newLayout) const {
+bool LLVMPointerType::areCompatible(
+    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
+    DataLayoutSpecInterface newSpec,
+    const DataLayoutIdentifiedEntryMap &map) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;
@@ -597,8 +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) const {
+bool LLVMStructType::areCompatible(
+    DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout,
+    DataLayoutSpecInterface newSpec,
+    const DataLayoutIdentifiedEntryMap &map) const {
   for (DataLayoutEntryInterface newEntry : newLayout) {
     if (!newEntry.isTypeEntry())
       continue;

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,22 @@ 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 DataLayoutIdentifiedEntryMap &map) 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 +66,12 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
           return false;
         });
     if (it == oldLayout.end()) {
+      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);
-          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,

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -246,6 +246,16 @@ 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)
+    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 +615,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)