address reviewer comments

Author: fabianmcg

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;

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; }]
     >,

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();
 }
 

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;

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

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