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 @@ -101,6 +101,9 @@ /// Returns the alloca memory space identifier. StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const; + /// Returns the stack alignment identifier. + StringAttr getStackAlignmentIdentifier(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 @@ -39,6 +39,9 @@ constexpr const static ::llvm::StringLiteral kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutStackAlignmentKey = "dlti.stack_alignment"; }]; 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 @@ -60,6 +60,10 @@ /// DataLayoutInterface if specified, otherwise returns the default. Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry); +/// Default handler for the stack alignment request. Dispatches to the +/// DataLayoutInterface if specified, otherwise returns the default. +unsigned getDefaultStackAlignment(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. @@ -166,6 +170,12 @@ /// Returns the memory space used for AllocaOps. Attribute getAllocaMemorySpace() 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 + /// unspecified. + unsigned getStackAlignment() const; + private: /// Combined layout spec at the given scope. const DataLayoutSpecInterface originalLayout; @@ -190,6 +200,9 @@ /// Cache for alloca memory space. mutable std::optional allocaMemorySpace; + + /// Cache for stack alignment. + mutable std::optional stackAlignment; }; } // 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 @@ -112,6 +112,12 @@ /*methodName=*/"getAllocaMemorySpaceIdentifier", /*args=*/(ins "::mlir::MLIRContext *":$context) >, + InterfaceMethod< + /*description=*/"Returns the stack alignment identifier.", + /*retTy=*/"::mlir::StringAttr", + /*methodName=*/"getStackAlignmentIdentifier", + /*args=*/(ins "::mlir::MLIRContext *":$context) + >, // Implementations may override this if they have an efficient lookup // mechanism. InterfaceMethod< @@ -274,6 +280,18 @@ return ::mlir::detail::getDefaultAllocaMemorySpace(entry); }] >, + StaticInterfaceMethod< + /*description=*/"Returns the natural stack alignment in bits computed " + "using the relevant entries. The data layout object " + "can be used for recursive queries.", + /*retTy=*/"unsigned", + /*methodName=*/"getStackAlignment", + /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultStackAlignment(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 @@ -108,6 +108,7 @@ constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword; constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey; +constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutStackAlignmentKey; namespace mlir { namespace impl { @@ -281,6 +282,12 @@ DLTIDialect::kDataLayoutAllocaMemorySpaceKey); } +StringAttr +DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr( + DLTIDialect::kDataLayoutStackAlignmentKey); +} + /// Parses an attribute with syntax /// attr ::= `#target.` `dl_spec` `<` attr-list? `>` /// attr-list ::= attr @@ -337,7 +344,8 @@ << DLTIDialect::kDataLayoutEndiannessBig << "' or '" << DLTIDialect::kDataLayoutEndiannessLittle << "'"; } - if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) + if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey || + 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 --- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp +++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp @@ -225,6 +225,17 @@ return entry.getValue(); } +// Returns the stack alignment if specified in the given entry. If the entry is +// empty the default alignment zero is returned. +unsigned +mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) { + if (entry == DataLayoutEntryInterface()) + return 0; + + auto attr = entry.getValue().cast(); + return attr.getValue().getZExtValue(); +} + DataLayoutEntryList mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries, TypeID typeID) { @@ -359,7 +370,7 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op) : originalLayout(getCombinedDataLayout(op)), scope(op), - allocaMemorySpace(std::nullopt) { + allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -368,7 +379,7 @@ mlir::DataLayout::DataLayout(ModuleOp op) : originalLayout(getCombinedDataLayout(op)), scope(op), - allocaMemorySpace(std::nullopt) { + allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) { #if LLVM_ENABLE_ABI_BREAKING_CHECKS checkMissingLayout(originalLayout, op); collectParentLayouts(op, layoutStack); @@ -486,6 +497,22 @@ return *allocaMemorySpace; } +unsigned mlir::DataLayout::getStackAlignment() const { + checkValid(); + MLIRContext *context = scope->getContext(); + if (stackAlignment) + return *stackAlignment; + DataLayoutEntryInterface entry; + if (originalLayout) + entry = originalLayout.getSpecForIdentifier( + originalLayout.getStackAlignmentIdentifier(context)); + if (auto iface = dyn_cast_or_null(scope)) + stackAlignment = iface.getStackAlignment(entry); + else + stackAlignment = detail::getDefaultStackAlignment(entry); + return *stackAlignment; +} + //===----------------------------------------------------------------------===// // DataLayoutSpecInterface //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h --- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h +++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h @@ -38,11 +38,11 @@ FloatType getFloatType(MLIRContext *context, unsigned width); /// Helper class that translates an LLVM data layout to an MLIR data layout -/// specification. Only integer, float, pointer, alloca memory space, and -/// endianness entries are translated. The class also returns all entries from -/// the default data layout specification found in the language reference -/// (https://llvm.org/docs/LangRef.html#data-layout) if they are not overwritten -/// by the provided data layout. +/// specification. Only integer, float, pointer, alloca memory space, stack +/// alignment, and endianness entries are translated. The class also returns all +/// entries from the default data layout specification found in the language +/// reference (https://llvm.org/docs/LangRef.html#data-layout) if they are not +/// overwritten by the provided data layout. class DataLayoutImporter { public: DataLayoutImporter(MLIRContext *context, @@ -99,6 +99,9 @@ /// Adds an alloca address space entry if there is none yet. LogicalResult tryToEmplaceAllocaAddrSpaceEntry(StringRef token); + /// Adds a stack alignment entry if there is none yet. + LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token); + std::string layoutStr = {}; StringRef lastToken = {}; SmallVector unhandledTokens; diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp --- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp +++ b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp @@ -183,6 +183,26 @@ return success(); } +LogicalResult +DataLayoutImporter::tryToEmplaceStackAlignmentEntry(StringRef token) { + auto key = + StringAttr::get(context, DLTIDialect::kDataLayoutStackAlignmentKey); + if (keyEntries.count(key)) + return success(); + + FailureOr alignment = tryToParseInt(token); + if (failed(alignment)) + return failure(); + + // Only store the stack alignment if it has a non-default value. + if (*alignment == 0) + return success(); + OpBuilder builder(context); + keyEntries.try_emplace(key, DataLayoutEntryAttr::get( + key, builder.getI32IntegerAttr(*alignment))); + return success(); +} + void DataLayoutImporter::translateDataLayout( const llvm::DataLayout &llvmDataLayout) { dataLayout = {}; @@ -230,6 +250,12 @@ return; continue; } + // Parse the stack alignment. + if (*prefix == "S") { + if (failed(tryToEmplaceStackAlignmentEntry(token))) + return; + continue; + } // Parse integer alignment specifications. if (*prefix == "i") { FailureOr width = tryToParseInt(token); 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 @@ -75,17 +75,28 @@ auto value = entry.getValue().cast(); bool isLittleEndian = value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle; - layoutStream << (isLittleEndian ? "e" : "E"); + layoutStream << "-" << (isLittleEndian ? "e" : "E"); 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(); - } + uint64_t space = value.getValue().getZExtValue(); + // Skip the default address space. + if (space == 0) + continue; + layoutStream << "-A" << space; + layoutStream.flush(); + continue; + } + if (key.getValue() == DLTIDialect::kDataLayoutStackAlignmentKey) { + auto value = entry.getValue().cast(); + uint64_t alignment = value.getValue().getZExtValue(); + // Skip the default stack alignment. + if (alignment == 0) + continue; + layoutStream << "-S" << alignment; + layoutStream.flush(); continue; } emitError(*loc) << "unsupported data layout key " << key; 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 @@ -8,18 +8,21 @@ // CHECK: bitsize = 64 // CHECK: preferred = 8 // 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: preferred = 8 // 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: preferred = 8 // CHECK: size = 8 + // CHECK: stack_alignment = 0 "test.data_layout_query"() : () -> !llvm.ptr<5> return } @@ -31,7 +34,8 @@ #dlti.dl_entry : vector<3xi32>>, #dlti.dl_entry, dense<[64, 64, 64]> : vector<3xi32>>, #dlti.dl_entry, dense<[32, 64, 64]> : vector<3xi32>>, - #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32> + #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>, + #dlti.dl_entry<"dlti.stack_alignment", 128 : i32> >} { // CHECK: @spec func.func @spec() { @@ -40,31 +44,36 @@ // CHECK: bitsize = 32 // CHECK: preferred = 8 // 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: preferred = 8 // 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: preferred = 8 // 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: preferred = 8 // 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: preferred = 8 // CHECK: size = 4 - "test.data_layout_query"() : () -> !llvm.ptr<4> + // CHECK: stack_alignment = 128 + "test.data_layout_query"() : () -> !llvm.ptr<4> return } } diff --git a/mlir/test/Target/LLVMIR/Import/data-layout.ll b/mlir/test/Target/LLVMIR/Import/data-layout.ll --- a/mlir/test/Target/LLVMIR/Import/data-layout.ll +++ b/mlir/test/Target/LLVMIR/Import/data-layout.ll @@ -28,6 +28,7 @@ ; CHECK-DAG: #dlti.dl_entry, dense<[32, 64, 64, 32]> : vector<4xi32>> ; CHECK-DAG: #dlti.dl_entry, dense<32> : vector<4xi32>> ; CHECK-DAG: #dlti.dl_entry, dense<64> : vector<4xi32>> +; CHECK-DAG: #dlti.dl_entry<"dlti.stack_alignment", 128 : i32> target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128" ; // ----- @@ -36,13 +37,15 @@ ; CHECK: #dlti.dl_spec< ; CHECK-DAG: #dlti.dl_entry<"dlti.endianness", "big"> ; CHECK-DAG: #dlti.dl_entry, dense<[16, 32, 64, 128]> : vector<4xi32>> +; CHECK-DAG: #dlti.dl_entry, dense<[16, 32, 64, 16]> : vector<4xi32>> ; CHECK-DAG: #dlti.dl_entry<"dlti.alloca_memory_space", 1 : ui32> ; CHECK-DAG: #dlti.dl_entry : vector<2xi32>> -target datalayout = "E-p270:16:32:64:128-A1-i64:64:128" +target datalayout = "A1-E-p270:16:32:64:128-p271:16:32:64-i64:64:128" ; // ----- ; CHECK: dlti.dl_spec = ; CHECK: #dlti.dl_spec< ; CHECK-NOT: #dlti.dl_entry<"dlti.alloca_memory_space" -target datalayout = "E-A0" +; CHECK-NOT: #dlti.dl_entry<"dlti.stack_alignment" +target datalayout = "A0-S0" diff --git a/mlir/test/Target/LLVMIR/data-layout.mlir b/mlir/test/Target/LLVMIR/data-layout.mlir --- a/mlir/test/Target/LLVMIR/data-layout.mlir +++ b/mlir/test/Target/LLVMIR/data-layout.mlir @@ -2,12 +2,16 @@ // CHECK: target datalayout // CHECK: E- +// CHECK: A4- +// CHECK: S128- // CHECK: i64:64:128 // CHECK: f80:128:256 // CHECK: p0:32:64:128 // CHECK: p1:32:32:32:64 module attributes {dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry<"dlti.endianness", "big">, +#dlti.dl_entry<"dlti.alloca_memory_space", 4 : ui32>, +#dlti.dl_entry<"dlti.stack_alignment", 128 : i32>, #dlti.dl_entry, #dlti.dl_entry : vector<2xi32>>, #dlti.dl_entry : vector<2xi32>>, @@ -21,6 +25,22 @@ // ----- +// CHECK: target datalayout +// CHECK: e +// CHECK-NOT: A0 +// CHECK-NOT: S0 +module attributes {dlti.dl_spec = #dlti.dl_spec< +#dlti.dl_entry<"dlti.endianness", "little">, +#dlti.dl_entry<"dlti.alloca_memory_space", 0 : ui32>, +#dlti.dl_entry<"dlti.stack_alignment", 0 : i32> +>} { + llvm.func @bar() { + llvm.return + } +} + +// ----- + // expected-error@below {{unsupported data layout for non-signless integer 'ui64'}} module attributes {dlti.dl_spec = #dlti.dl_spec< #dlti.dl_entry : vector<2xi32>>> 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 @@ -41,6 +41,7 @@ unsigned alignment = layout.getTypeABIAlignment(op.getType()); unsigned preferred = layout.getTypePreferredAlignment(op.getType()); Attribute allocaMemorySpace = layout.getAllocaMemorySpace(); + unsigned stackAlignment = layout.getStackAlignment(); op->setAttrs( {builder.getNamedAttr("size", builder.getIndexAttr(size)), builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)), @@ -49,9 +50,9 @@ builder.getNamedAttr("alloca_memory_space", allocaMemorySpace == Attribute() ? builder.getUI32IntegerAttr(0) - : allocaMemorySpace) - - }); + : allocaMemorySpace), + builder.getNamedAttr("stack_alignment", + builder.getIndexAttr(stackAlignment))}); }); } }; 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 @@ -24,6 +24,8 @@ constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; constexpr static llvm::StringLiteral kAllocaKeyName = "dltest.alloca_memory_space"; +constexpr static llvm::StringLiteral kStackAlignmentKeyName = + "dltest.stack_alignment"; /// Trivial array storage for the custom data layout spec attribute, just a list /// of entries. @@ -67,6 +69,9 @@ StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const { return Builder(context).getStringAttr(kAllocaKeyName); } + StringAttr getStackAlignmentIdentifier(MLIRContext *context) const { + return Builder(context).getStringAttr(kStackAlignmentKeyName); + } }; /// A type subject to data layout that exits the program if it is queried more @@ -276,6 +281,7 @@ EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getStackAlignment(), 0u); } TEST(DataLayout, NullSpec) { @@ -302,6 +308,7 @@ EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getStackAlignment(), 0u); } TEST(DataLayout, EmptySpec) { @@ -327,6 +334,7 @@ EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); + EXPECT_EQ(layout.getStackAlignment(), 0u); } TEST(DataLayout, SpecWithEntries) { @@ -334,7 +342,8 @@ "dltest.op_with_layout"() { dltest.layout = #dltest.spec< #dlti.dl_entry, #dlti.dl_entry, - #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32> + #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>, + #dlti.dl_entry<"dltest.stack_alignment", 128 : i32> > } : () -> () )MLIR"; @@ -365,6 +374,7 @@ EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5)); + EXPECT_EQ(layout.getStackAlignment(), 128u); } TEST(DataLayout, Caching) {