diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td index 007a417653e4c..cf3a72633587d 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td +++ b/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; diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td index 136fe8781e717..1a08bafda54ee 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td +++ b/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"; diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h index 0603bc075ab5e..50d85a8ccb1be 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h @@ -34,6 +34,8 @@ using DataLayoutEntryList = llvm::SmallVector; using DataLayoutEntryListRef = llvm::ArrayRef; using TargetDeviceSpecListRef = llvm::ArrayRef; using TargetDeviceSpecEntry = std::pair; +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 endianness; /// Cache for the mangling mode. mutable std::optional manglingMode; - /// Cache for alloca, global, and program memory spaces. + /// Cache for default, alloca, global, and program memory spaces. + mutable std::optional defaultMemorySpace; mutable std::optional allocaMemorySpace; mutable std::optional programMemorySpace; mutable std::optional globalMemorySpace; diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td index 70551c4c5b144..818b441b9a770 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td +++ b/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; }] >, diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp index 92efecc62acd5..206d9f43a44a4 100644 --- a/mlir/lib/Dialect/DLTI/DLTI.cpp +++ b/mlir/lib/Dialect/DLTI/DLTI.cpp @@ -305,25 +305,8 @@ combineOneSpec(DataLayoutSpecInterface spec, DenseMap 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(kvp.second.front().getKey()); - assert(&typeSample.getDialect() != - typeSample.getContext()->getLoadedDialect() && - "unexpected data layout entry for built-in type"); - - auto interface = cast(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(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(kvp.second.front().getKey()); + assert(&typeSample.getDialect() != + typeSample.getContext()->getLoadedDialect() && + "unexpected data layout entry for built-in type"); + + auto interface = cast(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 || diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp index 453b206de294e..8f39ede721c92 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp +++ b/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()[static_cast(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; diff --git a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp b/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp index 49c8ed977d50c..1fc424c05f6c7 100644 --- a/mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp +++ b/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(llvm::cast(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(entry.getKey())) { auto ptrTy = llvm::cast(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 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, diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp index 2d098e2db6a96..7afa6ac3a00e4 100644 --- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp +++ b/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(scope)) + defaultMemorySpace = iface.getDefaultMemorySpace(entry); + else + defaultMemorySpace = detail::getDefaultMemorySpace(entry); + return *defaultMemorySpace; +} + mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const { checkValid(); if (allocaMemorySpace) diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir index 2dfc289c93065..4e60f991f1bf5 100644 --- a/mlir/test/Dialect/LLVMIR/layout.mlir +++ b/mlir/test/Dialect/LLVMIR/layout.mlir @@ -6,6 +6,7 @@ module { // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: default_memory_space = 0 // CHECK: endianness = "" // CHECK: global_memory_space = 0 // CHECK: index = 64 @@ -18,6 +19,7 @@ module { // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: default_memory_space = 0 // CHECK: endianness = "" // CHECK: global_memory_space = 0 // CHECK: index = 64 @@ -30,6 +32,7 @@ module { // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: default_memory_space = 0 // CHECK: endianness = "" // CHECK: global_memory_space = 0 // CHECK: index = 64 @@ -50,6 +53,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry, dense<[64, 64, 64]> : vector<3xi64>>, #dlti.dl_entry, dense<[32, 64, 64, 24]> : vector<4xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, + #dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>, #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>, #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>, #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>, @@ -61,6 +65,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: default_memory_space = 7 // CHECK: endianness = "little" // CHECK: global_memory_space = 2 // CHECK: index = 32 @@ -73,6 +78,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: default_memory_space = 7 // CHECK: endianness = "little" // CHECK: global_memory_space = 2 // CHECK: index = 32 @@ -84,6 +90,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 64 + // CHECK: default_memory_space = 7 // CHECK: endianness = "little" // CHECK: global_memory_space = 2 // CHECK: index = 64 @@ -96,6 +103,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: default_memory_space = 7 // CHECK: endianness = "little" // CHECK: global_memory_space = 2 // CHECK: index = 24 diff --git a/mlir/test/Dialect/Ptr/layout.mlir b/mlir/test/Dialect/Ptr/layout.mlir index 73189a388942a..356c57402651f 100644 --- a/mlir/test/Dialect/Ptr/layout.mlir +++ b/mlir/test/Dialect/Ptr/layout.mlir @@ -4,16 +4,18 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry>, #dlti.dl_entry,#ptr.spec>, #dlti.dl_entry, #ptr.spec>, + #dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>, #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>, #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>, #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64> >} { - // CHECK: @spec + // CHECK-LABEL: @spec func.func @spec() { // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: default_memory_space = 7 // CHECK: global_memory_space = 2 // CHECK: index = 32 // CHECK: preferred = 8 @@ -21,19 +23,21 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: size = 4 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !ptr.ptr - // CHECK: alignment = 4 + // CHECK: alignment = 1 // CHECK: alloca_memory_space = 5 - // CHECK: bitsize = 32 + // CHECK: bitsize = 64 + // CHECK: default_memory_space = 7 // CHECK: global_memory_space = 2 - // CHECK: index = 32 - // CHECK: preferred = 8 + // CHECK: index = 64 + // CHECK: preferred = 1 // CHECK: program_memory_space = 3 - // CHECK: size = 4 + // CHECK: size = 8 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !ptr.ptr<3> // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 64 + // CHECK: default_memory_space = 7 // CHECK: global_memory_space = 2 // CHECK: index = 64 // CHECK: preferred = 8 @@ -44,6 +48,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: default_memory_space = 7 // CHECK: global_memory_space = 2 // CHECK: index = 24 // CHECK: preferred = 8 @@ -57,6 +62,36 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // ----- +module attributes { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, #ptr.spec>, + #dlti.dl_entry<"dlti.default_memory_space", 1 : ui64> +>} { + // CHECK-LABEL: @default_memory_space + func.func @default_memory_space() { + // CHECK: alignment = 4 + // CHECK: bitsize = 32 + // CHECK: index = 32 + // CHECK: preferred = 4 + // CHECK: size = 4 + "test.data_layout_query"() : () -> !ptr.ptr + // CHECK: alignment = 4 + // CHECK: bitsize = 32 + // CHECK: index = 32 + // CHECK: preferred = 4 + // CHECK: size = 4 + "test.data_layout_query"() : () -> !ptr.ptr<1> + // CHECK: alignment = 4 + // CHECK: bitsize = 32 + // CHECK: index = 32 + // CHECK: preferred = 4 + // CHECK: size = 4 + "test.data_layout_query"() : () -> !ptr.ptr<2> + return + } +} + +// ----- + // expected-error@+2 {{preferred alignment is expected to be at least as large as ABI alignment}} module attributes { dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry> diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp index eae563b48c119..41a2fe5e4ee75 100644 --- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp +++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp @@ -42,6 +42,7 @@ struct TestDataLayoutQuery uint64_t preferred = layout.getTypePreferredAlignment(op.getType()); uint64_t index = layout.getTypeIndexBitwidth(op.getType()).value_or(0); Attribute endianness = layout.getEndianness(); + Attribute defaultMemorySpace = layout.getDefaultMemorySpace(); Attribute allocaMemorySpace = layout.getAllocaMemorySpace(); Attribute manglingMode = layout.getManglingMode(); Attribute programMemorySpace = layout.getProgramMemorySpace(); @@ -69,6 +70,10 @@ struct TestDataLayoutQuery builder.getNamedAttr("endianness", endianness == Attribute() ? builder.getStringAttr("") : endianness), + builder.getNamedAttr("default_memory_space", + defaultMemorySpace == Attribute() + ? builder.getUI32IntegerAttr(0) + : defaultMemorySpace), builder.getNamedAttr("alloca_memory_space", allocaMemorySpace == Attribute() ? builder.getUI32IntegerAttr(0) diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp index 0c237440834ef..5c784dcee6e15 100644 --- a/mlir/test/lib/Dialect/Test/TestTypes.cpp +++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp @@ -284,7 +284,9 @@ TestTypeWithLayoutType::getIndexBitwidth(const DataLayout &dataLayout, } bool TestTypeWithLayoutType::areCompatible( - DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout) const { + DataLayoutEntryListRef oldLayout, DataLayoutEntryListRef newLayout, + DataLayoutSpecInterface newSpec, + const DataLayoutIdentifiedEntryMap &map) const { unsigned old = extractKind(oldLayout, "alignment"); return old == 1 || extractKind(newLayout, "alignment") <= old; } diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp index 3d87948c5b5af..fd81f3021aa9b 100644 --- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp +++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -23,6 +23,8 @@ using namespace mlir; namespace { constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; constexpr static llvm::StringLiteral kEndiannesKeyName = "dltest.endianness"; +constexpr static llvm::StringLiteral kDefaultKeyName = + "dltest.default_memory_space"; constexpr static llvm::StringLiteral kAllocaKeyName = "dltest.alloca_memory_space"; constexpr static llvm::StringLiteral kManglingModeKeyName = @@ -82,6 +84,9 @@ struct CustomDataLayoutSpec StringAttr getEndiannessIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr(kEndiannesKeyName); } + StringAttr getDefaultMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr(kDefaultKeyName); + } StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr(kAllocaKeyName); } @@ -205,6 +210,15 @@ struct SingleQueryType return Attribute(); } + Attribute getDefaultMemorySpace(DataLayoutEntryInterface entry) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return Attribute(); + } + Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) { static bool executed = false; if (executed) @@ -475,6 +489,7 @@ module {} EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); EXPECT_EQ(layout.getEndianness(), Attribute()); + EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute()); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); @@ -508,6 +523,7 @@ TEST(DataLayout, NullSpec) { EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u); EXPECT_EQ(layout.getEndianness(), Attribute()); + EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute()); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); @@ -549,6 +565,7 @@ TEST(DataLayout, EmptySpec) { EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u); EXPECT_EQ(layout.getEndianness(), Attribute()); + EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute()); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); @@ -572,6 +589,7 @@ TEST(DataLayout, SpecWithEntries) { #dlti.dl_entry, #dlti.dl_entry, #dlti.dl_entry<"dltest.endianness", "little">, + #dlti.dl_entry<"dltest.default_memory_space", 1 : i32>, #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>, #dlti.dl_entry<"dltest.program_memory_space", 3 : i32>, #dlti.dl_entry<"dltest.global_memory_space", 2 : i32>, @@ -609,6 +627,7 @@ TEST(DataLayout, SpecWithEntries) { EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); EXPECT_EQ(layout.getEndianness(), Builder(&ctx).getStringAttr("little")); + EXPECT_EQ(layout.getDefaultMemorySpace(), Builder(&ctx).getI32IntegerAttr(1)); EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5)); EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3)); EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2));