[Mlir-commits] [mlir] [mlir][DataLayout] Keep consistent input/output order (PR #135185)

[llvmlistbot at llvm.org]

https://github.com/darkbuck updated https://github.com/llvm/llvm-project/pull/135185

>From bc04089e9d880b8b8ca8182f4a8ec1dbc0f1ea21 Mon Sep 17 00:00:00 2001
From: Michael Liao <michael.hliao at gmail.com>
Date: Thu, 10 Apr 2025 10:07:35 -0400
Subject: [PATCH] [mlir][DataLayout] Keep consistent input/output order

- Use 'MapVector' instead of 'DenseMap' to keep a consistent order when
  importing/printing entries to prevent run-by-run differences.
---
 .../mlir/Interfaces/DataLayoutInterfaces.td    |  6 +++---
 mlir/lib/Dialect/DLTI/DLTI.cpp                 | 18 +++++++++---------
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp   |  8 ++++----
 mlir/lib/Target/LLVMIR/DataLayoutImporter.h    |  5 +++--
 4 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index fc701b20cb007..c5973d4252b0a 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -232,9 +232,9 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
     /// identifier they are associated with. Users are not expected to call this
     /// method directly.
     void bucketEntriesByType(
-        ::llvm::DenseMap<::mlir::TypeID, ::mlir::DataLayoutEntryList> &types,
-        ::llvm::DenseMap<::mlir::StringAttr,
-                         ::mlir::DataLayoutEntryInterface> &ids);
+        ::llvm::MapVector<::mlir::TypeID, ::mlir::DataLayoutEntryList> &types,
+        ::llvm::MapVector<::mlir::StringAttr,
+                          ::mlir::DataLayoutEntryInterface> &ids);
   }];
 }
 
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 206d9f43a44a4..00cbd32708f1e 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -14,8 +14,8 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
-#include "llvm/ADT/TypeSwitch.h"
 
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/MathExtras.h"
@@ -293,16 +293,16 @@ overwriteDuplicateEntries(SmallVectorImpl<DataLayoutEntryInterface> &oldEntries,
 /// Combines a data layout spec into the given lists of entries organized by
 /// type class and identifier, overwriting them if necessary. Fails to combine
 /// if the two entries with identical keys are not compatible.
-static LogicalResult
-combineOneSpec(DataLayoutSpecInterface spec,
-               DenseMap<TypeID, DataLayoutEntryList> &entriesForType,
-               DenseMap<StringAttr, DataLayoutEntryInterface> &entriesForID) {
+static LogicalResult combineOneSpec(
+    DataLayoutSpecInterface spec,
+    llvm::MapVector<TypeID, DataLayoutEntryList> &entriesForType,
+    llvm::MapVector<StringAttr, DataLayoutEntryInterface> &entriesForID) {
   // A missing spec should be fine.
   if (!spec)
     return success();
 
-  DenseMap<TypeID, DataLayoutEntryList> newEntriesForType;
-  DenseMap<StringAttr, DataLayoutEntryInterface> newEntriesForID;
+  llvm::MapVector<TypeID, DataLayoutEntryList> newEntriesForType;
+  llvm::MapVector<StringAttr, DataLayoutEntryInterface> newEntriesForID;
   spec.bucketEntriesByType(newEntriesForType, newEntriesForID);
 
   // Combine non-Type DL entries first so they are visible to the
@@ -362,8 +362,8 @@ DataLayoutSpecAttr::combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
     return {};
 
   // Combine all specs in order, with `this` being the last one.
-  DenseMap<TypeID, DataLayoutEntryList> entriesForType;
-  DenseMap<StringAttr, DataLayoutEntryInterface> entriesForID;
+  llvm::MapVector<TypeID, DataLayoutEntryList> entriesForType;
+  llvm::MapVector<StringAttr, DataLayoutEntryInterface> entriesForID;
   for (DataLayoutSpecInterface spec : specs)
     if (failed(combineOneSpec(spec, entriesForType, entriesForID)))
       return nullptr;
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index 7afa6ac3a00e4..9b9bb0c500380 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -736,8 +736,8 @@ std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
 //===----------------------------------------------------------------------===//
 
 void DataLayoutSpecInterface::bucketEntriesByType(
-    DenseMap<TypeID, DataLayoutEntryList> &types,
-    DenseMap<StringAttr, DataLayoutEntryInterface> &ids) {
+    llvm::MapVector<TypeID, DataLayoutEntryList> &types,
+    llvm::MapVector<StringAttr, DataLayoutEntryInterface> &ids) {
   for (DataLayoutEntryInterface entry : getEntries()) {
     if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey()))
       types[type.getTypeID()].push_back(entry);
@@ -755,8 +755,8 @@ LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec,
 
   // Second, dispatch verifications of entry groups to types or dialects they
   // are associated with.
-  DenseMap<TypeID, DataLayoutEntryList> types;
-  DenseMap<StringAttr, DataLayoutEntryInterface> ids;
+  llvm::MapVector<TypeID, DataLayoutEntryList> types;
+  llvm::MapVector<StringAttr, DataLayoutEntryInterface> ids;
   spec.bucketEntriesByType(types, ids);
 
   for (const auto &kvp : types) {
diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
index 206c8dd486adb..491ff65f17a41 100644
--- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
+++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
@@ -17,6 +17,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "llvm/ADT/MapVector.h"
 
 namespace llvm {
 class StringRef;
@@ -110,8 +111,8 @@ class DataLayoutImporter {
   std::string layoutStr = {};
   StringRef lastToken = {};
   SmallVector<StringRef> unhandledTokens;
-  DenseMap<StringAttr, DataLayoutEntryInterface> keyEntries;
-  DenseMap<TypeAttr, DataLayoutEntryInterface> typeEntries;
+  llvm::MapVector<StringAttr, DataLayoutEntryInterface> keyEntries;
+  llvm::MapVector<TypeAttr, DataLayoutEntryInterface> typeEntries;
   MLIRContext *context;
   DataLayoutSpecInterface dataLayout;
 };