From c1ed45a271145acbfad81d87706aeebf361809c3 Mon Sep 17 00:00:00 2001 From: agozillon Date: Tue, 9 Jan 2024 13:56:11 +0100 Subject: [PATCH] [mlir] Add global and program memory space handling to the data layout subsystem (#77367) This patch is based on a previous PR https://reviews.llvm.org/D144657 that added alloca address space handling to MLIR's DataLayout and DLTI interface. This patch aims to add identical features to import and access the global and program memory space through MLIR's DataLayout/DLTI system. --- mlir/include/mlir/Dialect/DLTI/DLTI.h | 6 ++ mlir/include/mlir/Dialect/DLTI/DLTIBase.td | 6 ++ .../mlir/Interfaces/DataLayoutInterfaces.h | 18 +++++- .../mlir/Interfaces/DataLayoutInterfaces.td | 36 +++++++++++ mlir/lib/Dialect/DLTI/DLTI.cpp | 18 ++++++ mlir/lib/Interfaces/DataLayoutInterfaces.cpp | 62 ++++++++++++++++++- mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp | 23 +++++-- mlir/lib/Target/LLVMIR/DataLayoutImporter.h | 3 +- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 20 ++++++ mlir/test/Dialect/LLVMIR/layout.mlir | 18 ++++++ .../lib/Dialect/DLTI/TestDataLayoutQuery.cpp | 10 +++ .../Interfaces/DataLayoutInterfacesTest.cpp | 38 ++++++++++++ 12 files changed, 250 insertions(+), 8 deletions(-) diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h index b9a8763eb4498..bf23aa2d48a80 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTI.h +++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h @@ -103,6 +103,12 @@ class DataLayoutSpecAttr /// Returns the alloca memory space identifier. StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const; + /// Returns the program memory space identifier. + StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const; + + /// Returns the global memory space identifier. + StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const; + /// Returns the stack alignment identifier. StringAttr getStackAlignmentIdentifier(MLIRContext *context) const; diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td index 7d519f4efcd22..3572a99fad874 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td +++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td @@ -39,6 +39,12 @@ def DLTI_Dialect : Dialect { constexpr const static ::llvm::StringLiteral kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutProgramMemorySpaceKey = "dlti.program_memory_space"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutGlobalMemorySpaceKey = "dlti.global_memory_space"; constexpr const static ::llvm::StringLiteral kDataLayoutStackAlignmentKey = "dlti.stack_alignment"; diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h index f4b9b95fb89f8..4a21f76dfc5d1 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h @@ -61,6 +61,14 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout, /// DataLayoutInterface if specified, otherwise returns the default. Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry); +/// Default handler for program memory space request. Dispatches to the +/// DataLayoutInterface if specified, otherwise returns the default. +Attribute getDefaultProgramMemorySpace(DataLayoutEntryInterface entry); + +/// Default handler for global memory space request. Dispatches to the +/// DataLayoutInterface if specified, otherwise returns the default. +Attribute getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry); + /// Default handler for the stack alignment request. Dispatches to the /// DataLayoutInterface if specified, otherwise returns the default. uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry); @@ -175,6 +183,12 @@ class DataLayout { /// Returns the memory space used for AllocaOps. Attribute getAllocaMemorySpace() const; + /// Returns the memory space used for program memory operations. + Attribute getProgramMemorySpace() const; + + /// Returns the memory space used for global operations. + Attribute getGlobalMemorySpace() const; + /// Returns the natural alignment of the stack in bits. Alignment promotion of /// stack variables should be limited to the natural stack alignment to /// prevent dynamic stack alignment. Returns zero if the stack alignment is @@ -203,8 +217,10 @@ class DataLayout { mutable DenseMap abiAlignments; mutable DenseMap preferredAlignments; - /// Cache for alloca memory space. + /// Cache for alloca, global, and program memory spaces. mutable std::optional allocaMemorySpace; + mutable std::optional programMemorySpace; + mutable std::optional globalMemorySpace; /// Cache for stack alignment. mutable std::optional stackAlignment; diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td index 2f60a16baf50b..a8def967fffcf 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td @@ -112,6 +112,18 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> { /*methodName=*/"getAllocaMemorySpaceIdentifier", /*args=*/(ins "::mlir::MLIRContext *":$context) >, + InterfaceMethod< + /*description=*/"Returns the program memory space identifier.", + /*retTy=*/"::mlir::StringAttr", + /*methodName=*/"getProgramMemorySpaceIdentifier", + /*args=*/(ins "::mlir::MLIRContext *":$context) + >, + InterfaceMethod< + /*description=*/"Returns the global memory space identifier.", + /*retTy=*/"::mlir::StringAttr", + /*methodName=*/"getGlobalMemorySpaceIdentifier", + /*args=*/(ins "::mlir::MLIRContext *":$context) + >, InterfaceMethod< /*description=*/"Returns the stack alignment identifier.", /*retTy=*/"::mlir::StringAttr", @@ -280,6 +292,30 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> { return ::mlir::detail::getDefaultAllocaMemorySpace(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=*/"getProgramMemorySpace", + /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultProgramMemorySpace(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=*/"getGlobalMemorySpace", + /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultGlobalMemorySpace(entry); + }] + >, StaticInterfaceMethod< /*description=*/"Returns the natural stack alignment in bits computed " "using the relevant entries. The data layout object " diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp index aba9e7db0a2fe..daef2349430d0 100644 --- a/mlir/lib/Dialect/DLTI/DLTI.cpp +++ b/mlir/lib/Dialect/DLTI/DLTI.cpp @@ -108,6 +108,11 @@ void DataLayoutEntryAttr::print(AsmPrinter &os) const { constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword; constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey; +constexpr const StringLiteral + mlir::DLTIDialect::kDataLayoutProgramMemorySpaceKey; +constexpr const StringLiteral + mlir::DLTIDialect::kDataLayoutGlobalMemorySpaceKey; + constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutStackAlignmentKey; namespace mlir { @@ -282,6 +287,17 @@ DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const { DLTIDialect::kDataLayoutAllocaMemorySpaceKey); } +StringAttr DataLayoutSpecAttr::getProgramMemorySpaceIdentifier( + MLIRContext *context) const { + return Builder(context).getStringAttr( + DLTIDialect::kDataLayoutProgramMemorySpaceKey); +} + +StringAttr +DataLayoutSpecAttr::getGlobalMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr( + DLTIDialect::kDataLayoutGlobalMemorySpaceKey); +} StringAttr DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr( @@ -345,6 +361,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface { << DLTIDialect::kDataLayoutEndiannessLittle << "'"; } if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey || + entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey || + entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey || entryName == DLTIDialect::kDataLayoutStackAlignmentKey) return success(); return emitError(loc) << "unknown data layout entry name: " << entryName; diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp index 1178417fd2a6c..65c41f44192a9 100644 --- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp +++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp @@ -230,6 +230,30 @@ mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) { return entry.getValue(); } +// Returns the memory space used for the program 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::getDefaultProgramMemorySpace(DataLayoutEntryInterface entry) { + if (entry == DataLayoutEntryInterface()) { + return Attribute(); + } + + return entry.getValue(); +} + +// Returns the memory space used for global the global 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::getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry) { + if (entry == DataLayoutEntryInterface()) { + return Attribute(); + } + + return entry.getValue(); +} + // Returns the stack alignment if specified in the given entry. If the entry is // empty the default alignment zero is returned. uint64_t @@ -382,7 +406,8 @@ mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {} mlir::DataLayout::DataLayout(DataLayoutOpInterface op) : originalLayout(getCombinedDataLayout(op)), scope(op), - allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) { + allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt), + globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -391,7 +416,8 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op) mlir::DataLayout::DataLayout(ModuleOp op) : originalLayout(getCombinedDataLayout(op)), scope(op), - allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) { + allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt), + globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -510,6 +536,38 @@ mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const { return *allocaMemorySpace; } +mlir::Attribute mlir::DataLayout::getProgramMemorySpace() const { + checkValid(); + if (programMemorySpace) + return *programMemorySpace; + DataLayoutEntryInterface entry; + if (originalLayout) + entry = originalLayout.getSpecForIdentifier( + originalLayout.getProgramMemorySpaceIdentifier( + originalLayout.getContext())); + if (auto iface = dyn_cast_or_null(scope)) + programMemorySpace = iface.getProgramMemorySpace(entry); + else + programMemorySpace = detail::getDefaultProgramMemorySpace(entry); + return *programMemorySpace; +} + +mlir::Attribute mlir::DataLayout::getGlobalMemorySpace() const { + checkValid(); + if (globalMemorySpace) + return *globalMemorySpace; + DataLayoutEntryInterface entry; + if (originalLayout) + entry = originalLayout.getSpecForIdentifier( + originalLayout.getGlobalMemorySpaceIdentifier( + originalLayout.getContext())); + if (auto iface = dyn_cast_or_null(scope)) + globalMemorySpace = iface.getGlobalMemorySpace(entry); + else + globalMemorySpace = detail::getDefaultGlobalMemorySpace(entry); + return *globalMemorySpace; +} + uint64_t mlir::DataLayout::getStackAlignment() const { checkValid(); if (stackAlignment) diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp index 95f3cc074b1dd..392f552f480cc 100644 --- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp +++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp @@ -164,9 +164,9 @@ DataLayoutImporter::tryToEmplaceEndiannessEntry(StringRef endianness, } LogicalResult -DataLayoutImporter::tryToEmplaceAllocaAddrSpaceEntry(StringRef token) { - auto key = - StringAttr::get(context, DLTIDialect::kDataLayoutAllocaMemorySpaceKey); +DataLayoutImporter::tryToEmplaceAddrSpaceEntry(StringRef token, + llvm::StringLiteral spaceKey) { + auto key = StringAttr::get(context, spaceKey); if (keyEntries.count(key)) return success(); @@ -247,9 +247,24 @@ void DataLayoutImporter::translateDataLayout( return; continue; } + // Parse the program address space. + if (*prefix == "P") { + if (failed(tryToEmplaceAddrSpaceEntry( + token, DLTIDialect::kDataLayoutProgramMemorySpaceKey))) + return; + continue; + } + // Parse the global address space. + if (*prefix == "G") { + if (failed(tryToEmplaceAddrSpaceEntry( + token, DLTIDialect::kDataLayoutGlobalMemorySpaceKey))) + return; + continue; + } // Parse the alloca address space. if (*prefix == "A") { - if (failed(tryToEmplaceAllocaAddrSpaceEntry(token))) + if (failed(tryToEmplaceAddrSpaceEntry( + token, DLTIDialect::kDataLayoutAllocaMemorySpaceKey))) return; continue; } diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h index 15f0f7ddf0705..59b60acd24be2 100644 --- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h +++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h @@ -97,7 +97,8 @@ class DataLayoutImporter { StringRef token); /// Adds an alloca address space entry if there is none yet. - LogicalResult tryToEmplaceAllocaAddrSpaceEntry(StringRef token); + LogicalResult tryToEmplaceAddrSpaceEntry(StringRef token, + llvm::StringLiteral spaceKey); /// Adds a stack alignment entry if there is none yet. LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token); diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp index 1722d74c08b62..ce46a194ea7d9 100644 --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -190,6 +190,26 @@ translateDataLayout(DataLayoutSpecInterface attribute, layoutStream.flush(); continue; } + if (key.getValue() == DLTIDialect::kDataLayoutProgramMemorySpaceKey) { + auto value = cast(entry.getValue()); + uint64_t space = value.getValue().getZExtValue(); + // Skip the default address space. + if (space == 0) + continue; + layoutStream << "-P" << space; + layoutStream.flush(); + continue; + } + if (key.getValue() == DLTIDialect::kDataLayoutGlobalMemorySpaceKey) { + auto value = cast(entry.getValue()); + uint64_t space = value.getValue().getZExtValue(); + // Skip the default address space. + if (space == 0) + continue; + layoutStream << "-G" << space; + layoutStream.flush(); + continue; + } if (key.getValue() == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) { auto value = cast(entry.getValue()); uint64_t space = value.getValue().getZExtValue(); diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir index 99d6178538448..2868e1740f861 100644 --- a/mlir/test/Dialect/LLVMIR/layout.mlir +++ b/mlir/test/Dialect/LLVMIR/layout.mlir @@ -6,21 +6,27 @@ module { // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: global_memory_space = 0 // CHECK: preferred = 8 + // CHECK: program_memory_space = 0 // CHECK: size = 8 // CHECK: stack_alignment = 0 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: global_memory_space = 0 // CHECK: preferred = 8 + // CHECK: program_memory_space = 0 // CHECK: size = 8 // CHECK: stack_alignment = 0 "test.data_layout_query"() : () -> !llvm.ptr<3> // CHECK: alignment = 8 // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 + // CHECK: global_memory_space = 0 // CHECK: preferred = 8 + // CHECK: program_memory_space = 0 // CHECK: size = 8 // CHECK: stack_alignment = 0 "test.data_layout_query"() : () -> !llvm.ptr<5> @@ -35,6 +41,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry, dense<[64, 64, 64]> : vector<3xi64>>, #dlti.dl_entry, dense<[32, 64, 64]> : vector<3xi64>>, #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 @@ -42,35 +50,45 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: global_memory_space = 2 // CHECK: preferred = 8 + // CHECK: program_memory_space = 3 // CHECK: size = 4 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: global_memory_space = 2 // CHECK: preferred = 8 + // CHECK: program_memory_space = 3 // CHECK: size = 4 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !llvm.ptr<3> // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 64 + // CHECK: global_memory_space = 2 // CHECK: preferred = 8 + // CHECK: program_memory_space = 3 // CHECK: size = 8 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !llvm.ptr<5> // CHECK: alignment = 4 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: global_memory_space = 2 // CHECK: preferred = 8 + // CHECK: program_memory_space = 3 // CHECK: size = 4 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !llvm.ptr<3> // CHECK: alignment = 8 // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 + // CHECK: global_memory_space = 2 // CHECK: preferred = 8 + // CHECK: program_memory_space = 3 // CHECK: size = 4 // CHECK: stack_alignment = 128 "test.data_layout_query"() : () -> !llvm.ptr<4> diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp index 7e3d3f6dc3f0e..740562e778302 100644 --- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp +++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp @@ -41,6 +41,8 @@ struct TestDataLayoutQuery unsigned alignment = layout.getTypeABIAlignment(op.getType()); unsigned preferred = layout.getTypePreferredAlignment(op.getType()); Attribute allocaMemorySpace = layout.getAllocaMemorySpace(); + Attribute programMemorySpace = layout.getProgramMemorySpace(); + Attribute globalMemorySpace = layout.getGlobalMemorySpace(); unsigned stackAlignment = layout.getStackAlignment(); op->setAttrs( {builder.getNamedAttr("size", builder.getIndexAttr(size)), @@ -51,6 +53,14 @@ struct TestDataLayoutQuery allocaMemorySpace == Attribute() ? builder.getUI32IntegerAttr(0) : allocaMemorySpace), + builder.getNamedAttr("program_memory_space", + programMemorySpace == Attribute() + ? builder.getUI32IntegerAttr(0) + : programMemorySpace), + builder.getNamedAttr("global_memory_space", + globalMemorySpace == Attribute() + ? builder.getUI32IntegerAttr(0) + : globalMemorySpace), builder.getNamedAttr("stack_alignment", builder.getIndexAttr(stackAlignment))}); }); diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp index 79599b8c48509..794e19710fadc 100644 --- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp +++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -24,6 +24,10 @@ namespace { constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; constexpr static llvm::StringLiteral kAllocaKeyName = "dltest.alloca_memory_space"; +constexpr static llvm::StringLiteral kProgramKeyName = + "dltest.program_memory_space"; +constexpr static llvm::StringLiteral kGlobalKeyName = + "dltest.global_memory_space"; constexpr static llvm::StringLiteral kStackAlignmentKeyName = "dltest.stack_alignment"; @@ -72,6 +76,12 @@ struct CustomDataLayoutSpec StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr(kAllocaKeyName); } + StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr(kProgramKeyName); + } + StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr(kGlobalKeyName); + } StringAttr getStackAlignmentIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr(kStackAlignmentKeyName); } @@ -128,6 +138,24 @@ struct SingleQueryType executed = true; return Attribute(); } + + Attribute getProgramMemorySpace(DataLayoutEntryInterface entry) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return Attribute(); + } + + Attribute getGlobalMemorySpace(DataLayoutEntryInterface entry) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return Attribute(); + } }; /// A types that is not subject to data layout. @@ -290,6 +318,8 @@ module {} EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); + EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); EXPECT_EQ(layout.getStackAlignment(), 0u); } @@ -317,6 +347,8 @@ TEST(DataLayout, NullSpec) { EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); + EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); EXPECT_EQ(layout.getStackAlignment(), 0u); } @@ -343,6 +375,8 @@ TEST(DataLayout, EmptySpec) { EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); + EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); EXPECT_EQ(layout.getStackAlignment(), 0u); } @@ -352,6 +386,8 @@ TEST(DataLayout, SpecWithEntries) { #dlti.dl_entry, #dlti.dl_entry, #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>, #dlti.dl_entry<"dltest.stack_alignment", 128 : i32> > } : () -> () )MLIR"; @@ -383,6 +419,8 @@ TEST(DataLayout, SpecWithEntries) { EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5)); + EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3)); + EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2)); EXPECT_EQ(layout.getStackAlignment(), 128u); }