diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h --- a/mlir/include/mlir/Dialect/DLTI/DLTI.h +++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h @@ -98,6 +98,9 @@ /// Returns the list of entries. DataLayoutEntryListRef getEntries() const; + /// Returns the alloca memory space identifier. + StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const; + /// Parses an instance of this attribute. static DataLayoutSpecAttr parse(AsmParser &parser); diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td --- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td +++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td @@ -23,9 +23,11 @@ }]; let extraClassDeclaration = [{ + // Top level attribute name. constexpr const static ::llvm::StringLiteral kDataLayoutAttrName = "dlti.dl_spec"; + // Constants used in entries. constexpr const static ::llvm::StringLiteral kDataLayoutEndiannessKey = "dlti.endianness"; @@ -34,6 +36,9 @@ constexpr const static ::llvm::StringLiteral kDataLayoutEndiannessLittle = "little"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space"; }]; let useDefaultAttributePrinterParser = 1; diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h @@ -56,6 +56,10 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout, ArrayRef params); +/// Default handler for alloca memory space request. Dispatches to the +/// DataLayoutInterface if specified, otherwise returns the default. +Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry); + /// Given a list of data layout entries, returns a new list containing the /// entries with keys having the given type ID, i.e. belonging to the same type /// class. @@ -159,6 +163,9 @@ /// Returns the preferred of the given type in the current scope. unsigned getTypePreferredAlignment(Type t) const; + /// Returns the memory space used for AllocaOps. + Attribute getAllocaMemorySpace() const; + private: /// Combined layout spec at the given scope. const DataLayoutSpecInterface originalLayout; @@ -180,6 +187,9 @@ mutable DenseMap bitsizes; mutable DenseMap abiAlignments; mutable DenseMap preferredAlignments; + + /// Cache for alloca memory space. + mutable std::optional allocaMemorySpace; }; } // namespace mlir diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td @@ -106,6 +106,12 @@ /*methodName=*/"getEntries", /*args=*/(ins) >, + InterfaceMethod< + /*description=*/"Returns the alloca memory space identifier.", + /*retTy=*/"::mlir::StringAttr", + /*methodName=*/"getAllocaMemorySpaceIdentifier", + /*args=*/(ins "::mlir::MLIRContext *":$context) + >, // Implementations may override this if they have an efficient lookup // mechanism. InterfaceMethod< @@ -256,6 +262,18 @@ params); }] >, + 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=*/"getAllocaMemorySpace", + /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultAllocaMemorySpace(entry); + }] + >, ]; let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }]; diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp --- a/mlir/lib/Dialect/DLTI/DLTI.cpp +++ b/mlir/lib/Dialect/DLTI/DLTI.cpp @@ -106,6 +106,8 @@ //===----------------------------------------------------------------------===// // constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword; +constexpr const StringLiteral + mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey; namespace mlir { namespace impl { @@ -273,6 +275,12 @@ return getImpl()->entries; } +StringAttr +DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr( + DLTIDialect::kDataLayoutAllocaMemorySpaceKey); +} + /// Parses an attribute with syntax /// attr ::= `#target.` `dl_spec` `<` attr-list? `>` /// attr-list ::= attr @@ -329,6 +337,8 @@ << DLTIDialect::kDataLayoutEndiannessBig << "' or '" << DLTIDialect::kDataLayoutEndiannessLittle << "'"; } + if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) + 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 --- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp +++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp @@ -213,6 +213,18 @@ reportMissingDataLayout(type); } +// Returns the memory space used for allocal operations 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::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) { + if (entry == DataLayoutEntryInterface()) { + return Attribute(); + } + + return entry.getValue(); +} + DataLayoutEntryList mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries, TypeID typeID) { @@ -346,7 +358,8 @@ mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {} mlir::DataLayout::DataLayout(DataLayoutOpInterface op) - : originalLayout(getCombinedDataLayout(op)), scope(op) { + : originalLayout(getCombinedDataLayout(op)), scope(op), + allocaMemorySpace(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -354,7 +367,8 @@ } mlir::DataLayout::DataLayout(ModuleOp op) - : originalLayout(getCombinedDataLayout(op)), scope(op) { + : originalLayout(getCombinedDataLayout(op)), scope(op), + allocaMemorySpace(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -456,6 +470,22 @@ }); } +mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const { + checkValid(); + MLIRContext *context = scope->getContext(); + if (allocaMemorySpace) + return *allocaMemorySpace; + DataLayoutEntryInterface entry; + if (originalLayout) + entry = originalLayout.getSpecForIdentifier( + originalLayout.getAllocaMemorySpaceIdentifier(context)); + if (auto iface = dyn_cast_or_null(scope)) + allocaMemorySpace = iface.getAllocaMemorySpace(entry); + else + allocaMemorySpace = detail::getDefaultAllocaMemorySpace(entry); + return *allocaMemorySpace; +} + //===----------------------------------------------------------------------===// // DataLayoutSpecInterface //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -20,6 +20,7 @@ #include "mlir/Dialect/DLTI/DLTI.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/Builders.h" #include "mlir/IR/Matchers.h" #include "mlir/Interfaces/DataLayoutInterfaces.h" #include "mlir/Tools/mlir-translate/Translation.h" @@ -179,9 +180,10 @@ // Remaining unhandled default layout defaults // e (little endian if not set) // p[n]:64:64:64 (non zero address spaces have 64-bit properties) + // Alloca address space defaults to 0. std::string append = "p:64:64:64-S0-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f16:16:16-f64:" - "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64"; + "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64-A0"; if (layoutstr.empty()) layoutstr = append; else @@ -227,6 +229,18 @@ StringAttr::get(context, DLTIDialect::kDataLayoutEndiannessKey), value); entries.emplace_back(entry); + } else if (symbol == 'A') { + unsigned addressSpace; + if (parameter.getAsInteger(/*Radix=*/10, addressSpace)) + return nullptr; + // Skip storing if generic address space is defined. + if (addressSpace != 0) { + auto entry = DataLayoutEntryAttr::get( + StringAttr::get(context, + DLTIDialect::kDataLayoutAllocaMemorySpaceKey), + mlir::Builder(context).getUI32IntegerAttr(addressSpace)); + entries.emplace_back(entry); + } } } diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -79,6 +79,15 @@ layoutStream.flush(); continue; } + if (key.getValue() == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) { + auto value = entry.getValue().cast(); + if (value != 0) { + // Only emit non-default address space. + layoutStream << "A" << value; + layoutStream.flush(); + } + continue; + } emitError(*loc) << "unsupported data layout key " << key; return failure(); } diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir --- a/mlir/test/Dialect/LLVMIR/layout.mlir +++ b/mlir/test/Dialect/LLVMIR/layout.mlir @@ -4,37 +4,44 @@ // CHECK: @no_spec func.func @no_spec() { // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr> // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 - // CHECK: bitsize = 64 + // CHECK: alloca_memory_space = 0 + // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr<5> @@ -47,47 +54,56 @@ module attributes { dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry, dense<[32, 32, 64]> : vector<3xi32>>, #dlti.dl_entry, dense<[64, 64, 64]> : vector<3xi32>>, - #dlti.dl_entry, dense<[32, 64, 64]> : vector<3xi32>> + #dlti.dl_entry, dense<[32, 64, 64]> : vector<3xi32>>, + #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32> >} { // CHECK: @spec func.func @spec() { // CHECK: alignment = 4 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr> // CHECK: alignment = 4 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_memory_space = 5 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 - // CHECK: bitsize = 32 + // CHECK: alloca_memory_space = 5 + // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr<3> // CHECK: alignment = 8 - // CHECK: bitsize = 32 + // CHECK: alloca_memory_space = 5 + // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "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 --- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp +++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp @@ -40,11 +40,18 @@ unsigned bitsize = layout.getTypeSizeInBits(op.getType()); unsigned alignment = layout.getTypeABIAlignment(op.getType()); unsigned preferred = layout.getTypePreferredAlignment(op.getType()); + Attribute allocaMemorySpace = layout.getAllocaMemorySpace(); op->setAttrs( {builder.getNamedAttr("size", builder.getIndexAttr(size)), builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)), builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)), - builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))}); + builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)), + builder.getNamedAttr("alloca_memory_space", + allocaMemorySpace == Attribute() + ? builder.getUI32IntegerAttr(0) + : allocaMemorySpace) + + }); }); } }; diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp --- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp +++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -22,6 +22,8 @@ namespace { constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; +constexpr static llvm::StringLiteral kAllocaKeyName = + "dltest.alloca_memory_space"; /// Trivial array storage for the custom data layout spec attribute, just a list /// of entries. @@ -62,6 +64,9 @@ } DataLayoutEntryListRef getEntries() const { return getImpl()->entries; } LogicalResult verifySpec(Location loc) { return success(); } + StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr(kAllocaKeyName); + } }; /// A type subject to data layout that exits the program if it is queried more @@ -104,6 +109,15 @@ executed = true; return 4; } + + Attribute getAllocaMemorySpace(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. @@ -260,6 +274,8 @@ EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u); EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u); EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); + + EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); } TEST(DataLayout, NullSpec) { @@ -275,6 +291,7 @@ auto op = cast(module->getBody()->getOperations().front()); DataLayout layout(op); + EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u); EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u); EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u); @@ -283,6 +300,8 @@ EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u); EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u); EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); + + EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); } TEST(DataLayout, EmptySpec) { @@ -306,13 +325,16 @@ EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u); EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u); EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); + + EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); } TEST(DataLayout, SpecWithEntries) { const char *ir = R"MLIR( "dltest.op_with_layout"() { dltest.layout = #dltest.spec< #dlti.dl_entry, - #dlti.dl_entry + #dlti.dl_entry, + #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32> > } : () -> () )MLIR"; @@ -341,6 +363,8 @@ EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u); EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u); EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); + + EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5)); } TEST(DataLayout, Caching) {