Index: mlir/include/mlir/Dialect/DLTI/DLTI.h =================================================================== --- mlir/include/mlir/Dialect/DLTI/DLTI.h +++ mlir/include/mlir/Dialect/DLTI/DLTI.h @@ -98,6 +98,9 @@ /// Returns the list of entries. DataLayoutEntryListRef getEntries() const; + /// Returns the alloca address space identifier + StringAttr getAllocaAddressSpaceIdentifier(MLIRContext *context) const; + /// Parses an instance of this attribute. static DataLayoutSpecAttr parse(AsmParser &parser); Index: mlir/include/mlir/Dialect/DLTI/DLTIBase.td =================================================================== --- mlir/include/mlir/Dialect/DLTI/DLTIBase.td +++ 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 + kDataLayoutAllocaAddressSpaceKey = "dlti.alloca_addrspace"; }]; let useDefaultAttributePrinterParser = 1; Index: mlir/include/mlir/Interfaces/DataLayoutInterfaces.h =================================================================== --- mlir/include/mlir/Interfaces/DataLayoutInterfaces.h +++ mlir/include/mlir/Interfaces/DataLayoutInterfaces.h @@ -56,6 +56,10 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout, ArrayRef params); +/// Default handler for alloca address space request. Dispatches to the +/// DataLayoutInterface if specified, otherwise returns the default vale 0. +unsigned getDefaultAllocaAddressSpace(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 address space used for AllocaOps + unsigned getAllocaAddressSpace() 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; + + /// Alloca address space. + mutable unsigned allocaAddressSpace; }; } // namespace mlir Index: mlir/include/mlir/Interfaces/DataLayoutInterfaces.td =================================================================== --- mlir/include/mlir/Interfaces/DataLayoutInterfaces.td +++ mlir/include/mlir/Interfaces/DataLayoutInterfaces.td @@ -106,6 +106,12 @@ /*methodName=*/"getEntries", /*args=*/(ins) >, + InterfaceMethod< + /*description=*/"Returns the alloca address space identifier.", + /*retTy=*/"::mlir::StringAttr", + /*methodName=*/"getAllocaAddressSpaceIdentifier", + /*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 address space used by the ABI computed " + "using the relevant entries. The data layout object " + "can be used for recursive queries.", + /*retTy=*/"unsigned", + /*methodName=*/"getAllocaAddressSpace", + /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultAllocaAddressSpace(entry); + }] + >, ]; let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }]; Index: mlir/lib/Dialect/DLTI/DLTI.cpp =================================================================== --- mlir/lib/Dialect/DLTI/DLTI.cpp +++ mlir/lib/Dialect/DLTI/DLTI.cpp @@ -106,6 +106,8 @@ //===----------------------------------------------------------------------===// // constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword; +constexpr const StringLiteral + mlir::DLTIDialect::kDataLayoutAllocaAddressSpaceKey; namespace mlir { namespace impl { @@ -273,6 +275,12 @@ return getImpl()->entries; } +StringAttr DataLayoutSpecAttr::getAllocaAddressSpaceIdentifier( + MLIRContext *context) const { + return Builder(context).getStringAttr( + DLTIDialect::kDataLayoutAllocaAddressSpaceKey); +} + /// Parses an attribute with syntax /// attr ::= `#target.` `dl_spec` `<` attr-list? `>` /// attr-list ::= attr @@ -329,6 +337,14 @@ << DLTIDialect::kDataLayoutEndiannessBig << "' or '" << DLTIDialect::kDataLayoutEndiannessLittle << "'"; } + if (entryName == DLTIDialect::kDataLayoutAllocaAddressSpaceKey) { + auto value = entry.getValue().dyn_cast(); + if (value && value.getType().isUnsignedInteger(32)) + return success(); + return emitError(loc) + << "'" << entryName + << "' data layout entry is expected to be a UI32 value"; + } return emitError(loc) << "unknown data layout entry name: " << entryName; } }; Index: mlir/lib/Interfaces/DataLayoutInterfaces.cpp =================================================================== --- mlir/lib/Interfaces/DataLayoutInterfaces.cpp +++ mlir/lib/Interfaces/DataLayoutInterfaces.cpp @@ -213,6 +213,18 @@ reportMissingDataLayout(type); } +// Returns the address space used for allocal operations if specified in the +// given entry. If the entry is empty the default address space 0 is returned. +unsigned +mlir::detail::getDefaultAllocaAddressSpace(DataLayoutEntryInterface entry) { + if (entry == DataLayoutEntryInterface()) { + return 0; + } + + auto attr = entry.getValue().cast(); + return attr.getValue().getZExtValue(); +} + 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), + allocaAddressSpace(UINT_MAX) { #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), + allocaAddressSpace(UINT_MAX) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -456,6 +470,21 @@ }); } +unsigned mlir::DataLayout::getAllocaAddressSpace() const { + checkValid(); + if (allocaAddressSpace != UINT_MAX) + return allocaAddressSpace; + DataLayoutEntryInterface entry; + if (originalLayout) + entry = originalLayout.getSpecForIdentifier( + originalLayout.getAllocaAddressSpaceIdentifier(scope->getContext())); + if (auto iface = dyn_cast_or_null(scope)) + allocaAddressSpace = iface.getAllocaAddressSpace(entry); + else + allocaAddressSpace = detail::getDefaultAllocaAddressSpace(entry); + return allocaAddressSpace; +} + //===----------------------------------------------------------------------===// // DataLayoutSpecInterface //===----------------------------------------------------------------------===// Index: mlir/lib/Target/LLVMIR/ModuleImport.cpp =================================================================== --- mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ 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,15 @@ StringAttr::get(context, DLTIDialect::kDataLayoutEndiannessKey), value); entries.emplace_back(entry); + } else if (symbol == 'A') { + unsigned addressSpace; + if (parameter.getAsInteger(/*Radix=*/10, addressSpace)) + return nullptr; + auto entry = DataLayoutEntryAttr::get( + StringAttr::get(context, + DLTIDialect::kDataLayoutAllocaAddressSpaceKey), + mlir::Builder(context).getUI32IntegerAttr(addressSpace)); + entries.emplace_back(entry); } } Index: mlir/lib/Target/LLVMIR/ModuleTranslation.cpp =================================================================== --- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -78,6 +78,15 @@ layoutStream.flush(); continue; } + if (key.getValue() == DLTIDialect::kDataLayoutAllocaAddressSpaceKey) { + 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(); } Index: mlir/test/Dialect/LLVMIR/layout.mlir =================================================================== --- mlir/test/Dialect/LLVMIR/layout.mlir +++ mlir/test/Dialect/LLVMIR/layout.mlir @@ -4,37 +4,44 @@ // CHECK: @no_spec func.func @no_spec() { // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr> // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 0 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 - // CHECK: bitsize = 64 + // CHECK: alloca_addrspace = 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_addrspace", 5 : ui32> >} { // CHECK: @spec func.func @spec() { // CHECK: alignment = 4 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr> // CHECK: alignment = 4 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 8 + // CHECK: alloca_addrspace = 5 // CHECK: bitsize = 64 // CHECK: preferred = 8 // CHECK: size = 8 "test.data_layout_query"() : () -> !llvm.ptr // CHECK: alignment = 4 - // CHECK: bitsize = 32 + // CHECK: alloca_addrspace = 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_addrspace = 5 + // CHECK: bitsize = 32 // CHECK: preferred = 8 // CHECK: size = 4 "test.data_layout_query"() : () -> !llvm.ptr<4> Index: mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp =================================================================== --- mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp +++ mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp @@ -40,11 +40,14 @@ unsigned bitsize = layout.getTypeSizeInBits(op.getType()); unsigned alignment = layout.getTypeABIAlignment(op.getType()); unsigned preferred = layout.getTypePreferredAlignment(op.getType()); + unsigned allocaAddrSpace = layout.getAllocaAddressSpace(); 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_addrspace", + builder.getUI32IntegerAttr(allocaAddrSpace))}); }); } }; Index: mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp =================================================================== --- mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp +++ mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -22,6 +22,7 @@ namespace { constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; +constexpr static llvm::StringLiteral kAllocaKeyName = "dltest.alloca_addrspace"; /// Trivial array storage for the custom data layout spec attribute, just a list /// of entries. @@ -62,6 +63,9 @@ } DataLayoutEntryListRef getEntries() const { return getImpl()->entries; } LogicalResult verifySpec(Location loc) { return success(); } + StringAttr getAllocaAddressSpaceIdentifier(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 +108,15 @@ executed = true; return 4; } + + unsigned getAllocaAddressSpace(DataLayoutEntryInterface entry) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return 0; + } }; /// A types that is not subject to data layout. @@ -260,6 +273,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.getAllocaAddressSpace(), 0u); } TEST(DataLayout, NullSpec) { @@ -275,6 +290,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 +299,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.getAllocaAddressSpace(), 0u); } TEST(DataLayout, EmptySpec) { @@ -306,6 +324,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.getAllocaAddressSpace(), 0u); } TEST(DataLayout, SpecWithEntries) { @@ -341,6 +361,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.getAllocaAddressSpace(), 0u); } TEST(DataLayout, Caching) {