Skip to content

Commit 9188288

Browse files
darkbuckdarkbuck
authored
[mlir][DataLayout] Keep consistent input/output order (#135185)
- Use 'MapVector' instead of 'DenseMap' to keep a consistent order when importing/printing entries to prevent run-by-run differences.
1 parent 53ae2bd commit 9188288

File tree

6 files changed

+46
-44
lines changed

6 files changed

+46
-44
lines changed

‎mlir/include/mlir/Interfaces/DataLayoutInterfaces.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "mlir/IR/DialectInterface.h"
2020
#include "mlir/IR/OpDefinition.h"
2121
#include "llvm/ADT/DenseMap.h"
22+
#include "llvm/ADT/MapVector.h"
2223
#include "llvm/Support/TypeSize.h"
2324

2425
namespace mlir {
@@ -35,7 +36,7 @@ using DataLayoutEntryListRef = llvm::ArrayRef<DataLayoutEntryInterface>;
3536
using TargetDeviceSpecListRef = llvm::ArrayRef<TargetDeviceSpecInterface>;
3637
using TargetDeviceSpecEntry = std::pair<StringAttr, TargetDeviceSpecInterface>;
3738
using DataLayoutIdentifiedEntryMap =
38-
::llvm::DenseMap<::mlir::StringAttr, ::mlir::DataLayoutEntryInterface>;
39+
::llvm::MapVector<::mlir::StringAttr, ::mlir::DataLayoutEntryInterface>;
3940
class DataLayoutOpInterface;
4041
class DataLayoutSpecInterface;
4142
class ModuleOp;

‎mlir/include/mlir/Interfaces/DataLayoutInterfaces.td‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -232,9 +232,9 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
232232
/// identifier they are associated with. Users are not expected to call this
233233
/// method directly.
234234
void bucketEntriesByType(
235-
::llvm::DenseMap<::mlir::TypeID, ::mlir::DataLayoutEntryList> &types,
236-
::llvm::DenseMap<::mlir::StringAttr,
237-
::mlir::DataLayoutEntryInterface> &ids);
235+
::llvm::MapVector<::mlir::TypeID, ::mlir::DataLayoutEntryList> &types,
236+
::llvm::MapVector<::mlir::StringAttr,
237+
::mlir::DataLayoutEntryInterface> &ids);
238238
}];
239239
}
240240

‎mlir/lib/Dialect/DLTI/DLTI.cpp‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
#include "mlir/IR/BuiltinTypes.h"
1515
#include "mlir/IR/Dialect.h"
1616
#include "mlir/IR/DialectImplementation.h"
17-
#include "llvm/ADT/TypeSwitch.h"
1817

18+
#include "llvm/ADT/MapVector.h"
1919
#include "llvm/ADT/TypeSwitch.h"
2020
#include "llvm/Support/Debug.h"
2121
#include "llvm/Support/MathExtras.h"
@@ -293,16 +293,16 @@ overwriteDuplicateEntries(SmallVectorImpl<DataLayoutEntryInterface> &oldEntries,
293293
/// Combines a data layout spec into the given lists of entries organized by
294294
/// type class and identifier, overwriting them if necessary. Fails to combine
295295
/// if the two entries with identical keys are not compatible.
296-
static LogicalResult
297-
combineOneSpec(DataLayoutSpecInterface spec,
298-
DenseMap<TypeID, DataLayoutEntryList> &entriesForType,
299-
DenseMap<StringAttr, DataLayoutEntryInterface> &entriesForID) {
296+
static LogicalResult combineOneSpec(
297+
DataLayoutSpecInterface spec,
298+
llvm::MapVector<TypeID, DataLayoutEntryList> &entriesForType,
299+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> &entriesForID) {
300300
// A missing spec should be fine.
301301
if (!spec)
302302
return success();
303303

304-
DenseMap<TypeID, DataLayoutEntryList> newEntriesForType;
305-
DenseMap<StringAttr, DataLayoutEntryInterface> newEntriesForID;
304+
llvm::MapVector<TypeID, DataLayoutEntryList> newEntriesForType;
305+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> newEntriesForID;
306306
spec.bucketEntriesByType(newEntriesForType, newEntriesForID);
307307

308308
// Combine non-Type DL entries first so they are visible to the
@@ -362,8 +362,8 @@ DataLayoutSpecAttr::combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
362362
return {};
363363

364364
// Combine all specs in order, with `this` being the last one.
365-
DenseMap<TypeID, DataLayoutEntryList> entriesForType;
366-
DenseMap<StringAttr, DataLayoutEntryInterface> entriesForID;
365+
llvm::MapVector<TypeID, DataLayoutEntryList> entriesForType;
366+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> entriesForID;
367367
for (DataLayoutSpecInterface spec : specs)
368368
if (failed(combineOneSpec(spec, entriesForType, entriesForID)))
369369
return nullptr;
@@ -374,7 +374,7 @@ DataLayoutSpecAttr::combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
374374
SmallVector<DataLayoutEntryInterface> entries;
375375
llvm::append_range(entries, llvm::make_second_range(entriesForID));
376376
for (const auto &kvp : entriesForType)
377-
llvm::append_range(entries, kvp.getSecond());
377+
llvm::append_range(entries, kvp.second);
378378

379379
return DataLayoutSpecAttr::get(getContext(), entries);
380380
}

‎mlir/lib/Interfaces/DataLayoutInterfaces.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -736,8 +736,8 @@ std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
736736
//===----------------------------------------------------------------------===//
737737

738738
void DataLayoutSpecInterface::bucketEntriesByType(
739-
DenseMap<TypeID, DataLayoutEntryList> &types,
740-
DenseMap<StringAttr, DataLayoutEntryInterface> &ids) {
739+
llvm::MapVector<TypeID, DataLayoutEntryList> &types,
740+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> &ids) {
741741
for (DataLayoutEntryInterface entry : getEntries()) {
742742
if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey()))
743743
types[type.getTypeID()].push_back(entry);
@@ -755,8 +755,8 @@ LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec,
755755

756756
// Second, dispatch verifications of entry groups to types or dialects they
757757
// are associated with.
758-
DenseMap<TypeID, DataLayoutEntryList> types;
759-
DenseMap<StringAttr, DataLayoutEntryInterface> ids;
758+
llvm::MapVector<TypeID, DataLayoutEntryList> types;
759+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> ids;
760760
spec.bucketEntriesByType(types, ids);
761761

762762
for (const auto &kvp : types) {

‎mlir/lib/Target/LLVMIR/DataLayoutImporter.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
1818
#include "mlir/IR/BuiltinAttributes.h"
1919
#include "mlir/Interfaces/DataLayoutInterfaces.h"
20+
#include "llvm/ADT/MapVector.h"
2021

2122
namespace llvm {
2223
class StringRef;
@@ -110,8 +111,8 @@ class DataLayoutImporter {
110111
std::string layoutStr = {};
111112
StringRef lastToken = {};
112113
SmallVector<StringRef> unhandledTokens;
113-
DenseMap<StringAttr, DataLayoutEntryInterface> keyEntries;
114-
DenseMap<TypeAttr, DataLayoutEntryInterface> typeEntries;
114+
llvm::MapVector<StringAttr, DataLayoutEntryInterface> keyEntries;
115+
llvm::MapVector<TypeAttr, DataLayoutEntryInterface> typeEntries;
115116
MLIRContext *context;
116117
DataLayoutSpecInterface dataLayout;
117118
};

‎mlir/test/Target/LLVMIR/Import/data-layout.ll‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -4,43 +4,43 @@
44

55
; CHECK: dlti.dl_spec =
66
; CHECK: #dlti.dl_spec<
7-
; CHECK-DAG: "dlti.endianness" = "little"
8-
; CHECK-DAG: i1 = dense<8> : vector<2xi64>
9-
; CHECK-DAG: i8 = dense<8> : vector<2xi64>
10-
; CHECK-DAG: i16 = dense<16> : vector<2xi64>
11-
; CHECK-DAG: i32 = dense<32> : vector<2xi64>
12-
; CHECK-DAG: i64 = dense<[32, 64]> : vector<2xi64>
13-
; CHECK-DAG: !llvm.ptr = dense<64> : vector<4xi64>
14-
; CHECK-DAG: f16 = dense<16> : vector<2xi64>
15-
; CHECK-DAG: f64 = dense<64> : vector<2xi64>
16-
; CHECK-DAG: f128 = dense<128> : vector<2xi64>
7+
; CHECK-SAME: !llvm.ptr = dense<64> : vector<4xi64>
8+
; CHECK-SAME: i1 = dense<8> : vector<2xi64>
9+
; CHECK-SAME: i8 = dense<8> : vector<2xi64>
10+
; CHECK-SAME: i16 = dense<16> : vector<2xi64>
11+
; CHECK-SAME: i32 = dense<32> : vector<2xi64>
12+
; CHECK-SAME: i64 = dense<[32, 64]> : vector<2xi64>
13+
; CHECK-SAME: f16 = dense<16> : vector<2xi64>
14+
; CHECK-SAME: f64 = dense<64> : vector<2xi64>
15+
; CHECK-SAME: f128 = dense<128> : vector<2xi64>
16+
; CHECK-SAME: "dlti.endianness" = "little"
1717
; CHECK: >
1818
target datalayout = ""
1919

2020
; // -----
2121

2222
; CHECK: dlti.dl_spec =
2323
; CHECK: #dlti.dl_spec<
24-
; CHECK-DAG: "dlti.endianness" = "little"
25-
; CHECK-DAG: i64 = dense<64> : vector<2xi64>
26-
; CHECK-DAG: f80 = dense<128> : vector<2xi64>
27-
; CHECK-DAG: i8 = dense<8> : vector<2xi64>
28-
; CHECK-DAG: !llvm.ptr<270> = dense<[32, 64, 64, 32]> : vector<4xi64>
29-
; CHECK-DAG: !llvm.ptr<271> = dense<32> : vector<4xi64>
30-
; CHECK-DAG: !llvm.ptr<272> = dense<64> : vector<4xi64>
31-
; CHECK-DAG: "dlti.stack_alignment" = 128 : i64
32-
; CHECK-DAG: "dlti.mangling_mode" = "e"
24+
; CHECK-SAME: !llvm.ptr<270> = dense<[32, 64, 64, 32]> : vector<4xi64>
25+
; CHECK-SAME: !llvm.ptr<271> = dense<32> : vector<4xi64>
26+
; CHECK-SAME: !llvm.ptr<272> = dense<64> : vector<4xi64>
27+
; CHECK-SAME: i64 = dense<64> : vector<2xi64>
28+
; CHECK-SAME: f80 = dense<128> : vector<2xi64>
29+
; CHECK-SAME: i8 = dense<8> : vector<2xi64>
30+
; CHECK-SAME: "dlti.endianness" = "little"
31+
; CHECK-SAME: "dlti.mangling_mode" = "e"
32+
; CHECK-SAME: "dlti.stack_alignment" = 128 : i64
3333
target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
3434

3535
; // -----
3636

3737
; CHECK: dlti.dl_spec =
3838
; CHECK: #dlti.dl_spec<
39-
; CHECK-DAG: "dlti.endianness" = "big"
40-
; CHECK-DAG: !llvm.ptr<270> = dense<[16, 32, 64, 8]> : vector<4xi64>
41-
; CHECK-DAG: !llvm.ptr<271> = dense<[16, 32, 64, 16]> : vector<4xi64>
42-
; CHECK-DAG: "dlti.alloca_memory_space" = 1 : ui64
43-
; CHECK-DAG: i64 = dense<[64, 128]> : vector<2xi64>
39+
; CHECK-SAME: !llvm.ptr<270> = dense<[16, 32, 64, 8]> : vector<4xi64>
40+
; CHECK-SAME: !llvm.ptr<271> = dense<[16, 32, 64, 16]> : vector<4xi64>
41+
; CHECK-SAME: i64 = dense<[64, 128]> : vector<2xi64>
42+
; CHECK-SAME: "dlti.alloca_memory_space" = 1 : ui64
43+
; CHECK-SAME: "dlti.endianness" = "big"
4444
target datalayout = "A1-E-p270:16:32:64:8-p271:16:32:64-i64:64:128"
4545

4646
; // -----

0 commit comments

Comments
 (0)