Skip to content

Commit 48375ac

Browse files
fabianmcgfabianmcg
committed
[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.
1 parent 96f3697 commit 48375ac

File tree

13 files changed

+178
-35
lines changed

13 files changed

+178
-35
lines changed

‎mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,9 @@ def DLTI_DataLayoutSpecAttr :
7474
/// Returns the endiannes identifier.
7575
StringAttr getEndiannessIdentifier(MLIRContext *context) const;
7676

77+
/// Returns the default memory space identifier.
78+
StringAttr getDefaultMemorySpaceIdentifier(MLIRContext *context) const;
79+
7780
/// Returns the alloca memory space identifier.
7881
StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
7982

‎mlir/include/mlir/Dialect/DLTI/DLTIBase.td‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ def DLTI_Dialect : Dialect {
5353
constexpr const static ::llvm::StringLiteral
5454
kDataLayoutManglingModeKey = "dlti.mangling_mode";
5555

56+
constexpr const static ::llvm::StringLiteral
57+
kDataLayoutDefaultMemorySpaceKey = "dlti.default_memory_space";
58+
5659
constexpr const static ::llvm::StringLiteral
5760
kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
5861

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,10 @@ getDefaultIndexBitwidth(Type type, const DataLayout &dataLayout,
7474
/// DataLayoutInterface if specified, otherwise returns the default.
7575
Attribute getDefaultEndianness(DataLayoutEntryInterface entry);
7676

77+
/// Default handler for the default memory space request. Dispatches to the
78+
/// DataLayoutInterface if specified, otherwise returns the default.
79+
Attribute getDefaultMemorySpace(DataLayoutEntryInterface entry);
80+
7781
/// Default handler for alloca memory space request. Dispatches to the
7882
/// DataLayoutInterface if specified, otherwise returns the default.
7983
Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
@@ -231,6 +235,9 @@ class DataLayout {
231235
/// Returns the specified endianness.
232236
Attribute getEndianness() const;
233237

238+
/// Returns the default memory space used for memory operations.
239+
Attribute getDefaultMemorySpace() const;
240+
234241
/// Returns the memory space used for AllocaOps.
235242
Attribute getAllocaMemorySpace() const;
236243

@@ -285,7 +292,8 @@ class DataLayout {
285292
mutable std::optional<Attribute> endianness;
286293
/// Cache for the mangling mode.
287294
mutable std::optional<Attribute> manglingMode;
288-
/// Cache for alloca, global, and program memory spaces.
295+
/// Cache for default, alloca, global, and program memory spaces.
296+
mutable std::optional<Attribute> defaultMemorySpace;
289297
mutable std::optional<Attribute> allocaMemorySpace;
290298
mutable std::optional<Attribute> programMemorySpace;
291299
mutable std::optional<Attribute> globalMemorySpace;

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

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
135135
/*methodName=*/"getEndiannessIdentifier",
136136
/*args=*/(ins "::mlir::MLIRContext *":$context)
137137
>,
138+
InterfaceMethod<
139+
/*description=*/"Returns the default memory space identifier.",
140+
/*retTy=*/"::mlir::StringAttr",
141+
/*methodName=*/"getDefaultMemorySpaceIdentifier",
142+
/*args=*/(ins "::mlir::MLIRContext *":$context)
143+
>,
138144
InterfaceMethod<
139145
/*description=*/"Returns the alloca memory space identifier.",
140146
/*retTy=*/"::mlir::StringAttr",
@@ -475,6 +481,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
475481
return ::mlir::detail::getDefaultEndianness(entry);
476482
}]
477483
>,
484+
StaticInterfaceMethod<
485+
/*description=*/"Returns the memory space used by the ABI computed "
486+
"using the relevant entries. The data layout object "
487+
"can be used for recursive queries.",
488+
/*retTy=*/"::mlir::Attribute",
489+
/*methodName=*/"getDefaultMemorySpace",
490+
/*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
491+
/*methodBody=*/"",
492+
/*defaultImplementation=*/[{
493+
return ::mlir::detail::getDefaultMemorySpace(entry);
494+
}]
495+
>,
478496
StaticInterfaceMethod<
479497
/*description=*/"Returns the memory space used by the ABI computed "
480498
"using the relevant entries. The data layout object "
@@ -637,11 +655,13 @@ def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> {
637655
InterfaceMethod<
638656
/*desc=*/"Returns true if the two lists of entries are compatible, that "
639657
"is, that `newLayout` spec entries can be nested in an op with "
640-
"`oldLayout` spec entries.",
658+
"`oldLayout` spec entries. `newSpec` is provided to further"
659+
"query data from the spec, e.g. the default address space.",
641660
/*retTy=*/"bool",
642661
/*methodName=*/"areCompatible",
643662
/*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout,
644-
"::mlir::DataLayoutEntryListRef":$newLayout),
663+
"::mlir::DataLayoutEntryListRef":$newLayout,
664+
"::mlir::DataLayoutSpecInterface":$newSpec),
645665
/*methodBody=*/"",
646666
/*defaultImplementation=*/[{ return true; }]
647667
>,

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -318,7 +318,8 @@ combineOneSpec(DataLayoutSpecInterface spec,
318318
"unexpected data layout entry for built-in type");
319319

320320
auto interface = cast<DataLayoutTypeInterface>(typeSample);
321-
if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second))
321+
if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second,
322+
spec))
322323
return failure();
323324

324325
overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
@@ -379,6 +380,12 @@ DataLayoutSpecAttr::getEndiannessIdentifier(MLIRContext *context) const {
379380
return Builder(context).getStringAttr(DLTIDialect::kDataLayoutEndiannessKey);
380381
}
381382

383+
StringAttr DataLayoutSpecAttr::getDefaultMemorySpaceIdentifier(
384+
MLIRContext *context) const {
385+
return Builder(context).getStringAttr(
386+
DLTIDialect::kDataLayoutDefaultMemorySpaceKey);
387+
}
388+
382389
StringAttr
383390
DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
384391
return Builder(context).getStringAttr(
@@ -609,7 +616,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
609616
<< DLTIDialect::kDataLayoutEndiannessBig << "' or '"
610617
<< DLTIDialect::kDataLayoutEndiannessLittle << "'";
611618
}
612-
if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
619+
if (entryName == DLTIDialect::kDataLayoutDefaultMemorySpaceKey ||
620+
entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
613621
entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey ||
614622
entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
615623
entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -350,7 +350,8 @@ LLVMPointerType::getIndexBitwidth(const DataLayout &dataLayout,
350350
}
351351

352352
bool LLVMPointerType::areCompatible(DataLayoutEntryListRef oldLayout,
353-
DataLayoutEntryListRef newLayout) const {
353+
DataLayoutEntryListRef newLayout,
354+
DataLayoutSpecInterface newSpec) const {
354355
for (DataLayoutEntryInterface newEntry : newLayout) {
355356
if (!newEntry.isTypeEntry())
356357
continue;
@@ -598,7 +599,8 @@ static uint64_t extractStructSpecValue(Attribute attr, StructDLEntryPos pos) {
598599
}
599600

600601
bool LLVMStructType::areCompatible(DataLayoutEntryListRef oldLayout,
601-
DataLayoutEntryListRef newLayout) const {
602+
DataLayoutEntryListRef newLayout,
603+
DataLayoutSpecInterface newSpec) const {
602604
for (DataLayoutEntryInterface newEntry : newLayout) {
603605
if (!newEntry.isTypeEntry())
604606
continue;

‎mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,12 @@ using namespace mlir::ptr;
2323

2424
constexpr const static unsigned kDefaultPointerSizeBits = 64;
2525
constexpr const static unsigned kBitsInByte = 8;
26-
constexpr const static unsigned kDefaultPointerAlignment = 8;
27-
28-
static Attribute getDefaultMemorySpace(PtrType ptr) { return nullptr; }
26+
constexpr const static unsigned kDefaultPointerAlignmentBits = 8;
2927

3028
/// Searches the data layout for the pointer spec, returns nullptr if it is not
3129
/// found.
32-
static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type) {
30+
static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type,
31+
Attribute defaultMemorySpace) {
3332
for (DataLayoutEntryInterface entry : params) {
3433
if (!entry.isTypeEntry())
3534
continue;
@@ -41,20 +40,21 @@ static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type) {
4140
}
4241
// If not found, and this is the pointer to the default memory space, assume
4342
// 64-bit pointers.
44-
if (type.getMemorySpace() == getDefaultMemorySpace(type))
43+
if (type.getMemorySpace() == defaultMemorySpace)
4544
return SpecAttr::get(type.getContext(), kDefaultPointerSizeBits,
46-
kDefaultPointerAlignment, kDefaultPointerAlignment,
47-
kDefaultPointerSizeBits);
45+
kDefaultPointerAlignmentBits,
46+
kDefaultPointerAlignmentBits, kDefaultPointerSizeBits);
4847
return nullptr;
4948
}
5049

5150
bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
52-
DataLayoutEntryListRef newLayout) const {
51+
DataLayoutEntryListRef newLayout,
52+
DataLayoutSpecInterface newSpec) const {
5353
for (DataLayoutEntryInterface newEntry : newLayout) {
5454
if (!newEntry.isTypeEntry())
5555
continue;
5656
uint32_t size = kDefaultPointerSizeBits;
57-
uint32_t abi = kDefaultPointerAlignment;
57+
uint32_t abi = kDefaultPointerAlignmentBits;
5858
auto newType = llvm::cast<PtrType>(llvm::cast<Type>(newEntry.getKey()));
5959
const auto *it =
6060
llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
@@ -65,10 +65,13 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
6565
return false;
6666
});
6767
if (it == oldLayout.end()) {
68+
Attribute defaultMemorySpace =
69+
mlir::detail::getDefaultMemorySpace(newSpec.getSpecForIdentifier(
70+
newSpec.getDefaultMemorySpaceIdentifier(getContext())));
6871
it = llvm::find_if(oldLayout, [&](DataLayoutEntryInterface entry) {
6972
if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey())) {
7073
auto ptrTy = llvm::cast<PtrType>(type);
71-
return ptrTy.getMemorySpace() == getDefaultMemorySpace(ptrTy);
74+
return ptrTy.getMemorySpace() == defaultMemorySpace;
7275
}
7376
return false;
7477
});
@@ -90,43 +93,44 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
9093

9194
uint64_t PtrType::getABIAlignment(const DataLayout &dataLayout,
9295
DataLayoutEntryListRef params) const {
93-
if (SpecAttr spec = getPointerSpec(params, *this))
96+
Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
97+
if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
9498
return spec.getAbi() / kBitsInByte;
9599

96-
return dataLayout.getTypeABIAlignment(
97-
get(getContext(), getDefaultMemorySpace(*this)));
100+
return dataLayout.getTypeABIAlignment(get(getContext(), defaultMemorySpace));
98101
}
99102

100103
std::optional<uint64_t>
101104
PtrType::getIndexBitwidth(const DataLayout &dataLayout,
102105
DataLayoutEntryListRef params) const {
103-
if (SpecAttr spec = getPointerSpec(params, *this)) {
106+
Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
107+
if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace)) {
104108
return spec.getIndex() == SpecAttr::kOptionalSpecValue ? spec.getSize()
105109
: spec.getIndex();
106110
}
107111

108-
return dataLayout.getTypeIndexBitwidth(
109-
get(getContext(), getDefaultMemorySpace(*this)));
112+
return dataLayout.getTypeIndexBitwidth(get(getContext(), defaultMemorySpace));
110113
}
111114

112115
llvm::TypeSize PtrType::getTypeSizeInBits(const DataLayout &dataLayout,
113116
DataLayoutEntryListRef params) const {
114-
if (SpecAttr spec = getPointerSpec(params, *this))
117+
Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
118+
if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
115119
return llvm::TypeSize::getFixed(spec.getSize());
116120

117121
// For other memory spaces, use the size of the pointer to the default memory
118122
// space.
119-
return dataLayout.getTypeSizeInBits(
120-
get(getContext(), getDefaultMemorySpace(*this)));
123+
return dataLayout.getTypeSizeInBits(get(getContext(), defaultMemorySpace));
121124
}
122125

123126
uint64_t PtrType::getPreferredAlignment(const DataLayout &dataLayout,
124127
DataLayoutEntryListRef params) const {
125-
if (SpecAttr spec = getPointerSpec(params, *this))
128+
Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
129+
if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
126130
return spec.getPreferred() / kBitsInByte;
127131

128132
return dataLayout.getTypePreferredAlignment(
129-
get(getContext(), getDefaultMemorySpace(*this)));
133+
get(getContext(), defaultMemorySpace));
130134
}
131135

132136
LogicalResult PtrType::verifyEntries(DataLayoutEntryListRef entries,

‎mlir/lib/Interfaces/DataLayoutInterfaces.cpp‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,17 @@ Attribute mlir::detail::getDefaultEndianness(DataLayoutEntryInterface entry) {
246246
return entry.getValue();
247247
}
248248

249+
// Returns the default memory space if specified in the given entry. If the
250+
// entry is empty the default memory space represented by an empty attribute is
251+
// returned.
252+
Attribute mlir::detail::getDefaultMemorySpace(DataLayoutEntryInterface entry) {
253+
if (entry == DataLayoutEntryInterface()) {
254+
return Attribute();
255+
}
256+
257+
return entry.getValue();
258+
}
259+
249260
// Returns the memory space used for alloca operations if specified in the
250261
// given entry. If the entry is empty the default memory space represented by
251262
// an empty attribute is returned.
@@ -605,6 +616,22 @@ mlir::Attribute mlir::DataLayout::getEndianness() const {
605616
return *endianness;
606617
}
607618

619+
mlir::Attribute mlir::DataLayout::getDefaultMemorySpace() const {
620+
checkValid();
621+
if (defaultMemorySpace)
622+
return *defaultMemorySpace;
623+
DataLayoutEntryInterface entry;
624+
if (originalLayout)
625+
entry = originalLayout.getSpecForIdentifier(
626+
originalLayout.getDefaultMemorySpaceIdentifier(
627+
originalLayout.getContext()));
628+
if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
629+
defaultMemorySpace = iface.getDefaultMemorySpace(entry);
630+
else
631+
defaultMemorySpace = detail::getDefaultMemorySpace(entry);
632+
return *defaultMemorySpace;
633+
}
634+
608635
mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
609636
checkValid();
610637
if (allocaMemorySpace)

‎mlir/test/Dialect/LLVMIR/layout.mlir‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ module {
66
// CHECK: alignment = 8
77
// CHECK: alloca_memory_space = 0
88
// CHECK: bitsize = 64
9+
// CHECK: default_memory_space = 0
910
// CHECK: endianness = ""
1011
// CHECK: global_memory_space = 0
1112
// CHECK: index = 64
@@ -18,6 +19,7 @@ module {
1819
// CHECK: alignment = 8
1920
// CHECK: alloca_memory_space = 0
2021
// CHECK: bitsize = 64
22+
// CHECK: default_memory_space = 0
2123
// CHECK: endianness = ""
2224
// CHECK: global_memory_space = 0
2325
// CHECK: index = 64
@@ -30,6 +32,7 @@ module {
3032
// CHECK: alignment = 8
3133
// CHECK: alloca_memory_space = 0
3234
// CHECK: bitsize = 64
35+
// CHECK: default_memory_space = 0
3336
// CHECK: endianness = ""
3437
// CHECK: global_memory_space = 0
3538
// CHECK: index = 64
@@ -50,6 +53,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
5053
#dlti.dl_entry<!llvm.ptr<5>, dense<[64, 64, 64]> : vector<3xi64>>,
5154
#dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64, 24]> : vector<4xi64>>,
5255
#dlti.dl_entry<"dlti.endianness", "little">,
56+
#dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>,
5357
#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
5458
#dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
5559
#dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
@@ -61,6 +65,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
6165
// CHECK: alignment = 4
6266
// CHECK: alloca_memory_space = 5
6367
// CHECK: bitsize = 32
68+
// CHECK: default_memory_space = 7
6469
// CHECK: endianness = "little"
6570
// CHECK: global_memory_space = 2
6671
// CHECK: index = 32
@@ -73,6 +78,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
7378
// CHECK: alignment = 4
7479
// CHECK: alloca_memory_space = 5
7580
// CHECK: bitsize = 32
81+
// CHECK: default_memory_space = 7
7682
// CHECK: endianness = "little"
7783
// CHECK: global_memory_space = 2
7884
// CHECK: index = 32
@@ -84,6 +90,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
8490
// CHECK: alignment = 8
8591
// CHECK: alloca_memory_space = 5
8692
// CHECK: bitsize = 64
93+
// CHECK: default_memory_space = 7
8794
// CHECK: endianness = "little"
8895
// CHECK: global_memory_space = 2
8996
// CHECK: index = 64
@@ -96,6 +103,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
96103
// CHECK: alignment = 8
97104
// CHECK: alloca_memory_space = 5
98105
// CHECK: bitsize = 32
106+
// CHECK: default_memory_space = 7
99107
// CHECK: endianness = "little"
100108
// CHECK: global_memory_space = 2
101109
// CHECK: index = 24

0 commit comments

Comments
 (0)