[mlir] Add stack alignment to the data layout dialect.

[gysit]

Summary

The revision adds the stack alignment to the
data layout dialect and it extends the LLVM dialect
import and export to support the new data layout
entry.

One possible use case for the flag is the LLVM dialect
inliner. The LLVM inliner queries the flag to
determine if it is safe to update the alignment of an
existing alloca. We may want to perform the same
optimization inside of MLIR.

[Dinistro]

LGMT!

mlir/test/Dialect/LLVMIR/layout.mlir

76

Ultra nit: As you are already touching this file, could you indent this line correctly?

[gysit]

Address review comment.

mlir/include/mlir/Dialect/DLTI/DLTI.h

public:
   DataLayoutSpecAttr combineWith(ArrayRef<DataLayoutSpecInterface> specs) const;
 
   /// Returns the list of entries.
   DataLayoutEntryListRef getEntries() const;
 
   /// 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);
 
   /// Prints this attribute.
   void print(AsmPrinter &os) const;
 };
 
 } // namespace mlir
 
 #include "mlir/Dialect/DLTI/DLTIDialect.h.inc"
 
 #endif // MLIR_DIALECT_DLTI_DLTI_H

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

let extraClassDeclaration = [{
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessBig = "big";
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessLittle = "little";
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
   }];
 
   let useDefaultAttributePrinterParser = 1;
 }
 
 def DLTI_DataLayoutEntryAttr : DialectAttr<
     DLTI_Dialect,
     CPred<"$_self.isa<::mlir::DataLayoutEntryAttr>()">,

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

 unsigned
 getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
                              ArrayRef<DataLayoutEntryInterface> params);
 
 /// Default handler for alloca memory space request. Dispatches to the
 /// 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.
 DataLayoutEntryList filterEntriesForType(DataLayoutEntryListRef entries,
                                          TypeID typeID);
 
 /// Given a list of data layout entries, returns the entry that has the given
 /// identifier as key, if such an entry exists in the list.
public:
   unsigned getTypeABIAlignment(Type t) const;
 
   /// 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;
 
+  /// 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;
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   /// List of enclosing layout specs.
   SmallVector<DataLayoutSpecInterface, 2> layoutStack;
 #endif
 
   /// Asserts that the cache is still valid. Expensive in debug mode. No-op in
   /// release mode.
   void checkValid() const;
 
   /// Operation defining the scope of requests.
   Operation *scope;
 
   /// Caches for individual requests.
   mutable DenseMap<Type, unsigned> sizes;
   mutable DenseMap<Type, unsigned> bitsizes;
   mutable DenseMap<Type, unsigned> abiAlignments;
   mutable DenseMap<Type, unsigned> preferredAlignments;
 
   /// Cache for alloca memory space.
   mutable std::optional<Attribute> allocaMemorySpace;
+
+  /// Cache for stack alignment.
+  mutable std::optional<unsigned> stackAlignment;
 };
 
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

InterfaceMethod<
       /*args=*/(ins)
     >,
     InterfaceMethod<
       /*description=*/"Returns the alloca memory space identifier.",
       /*retTy=*/"::mlir::StringAttr",
       /*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<
       /*description=*/"Returns a copy of the entries related to a specific "
                       "type class regardles of type parameters. ",
       /*retTy=*/"::mlir::DataLayoutEntryList",
       /*methodName=*/"getSpecForType",
       /*args=*/(ins "::mlir::TypeID":$type),
StaticInterfaceMethod<
       /*retTy=*/"::mlir::Attribute",
       /*methodName=*/"getAllocaMemorySpace",
       /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         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); }];
 }
 
 //===----------------------------------------------------------------------===//
 // Type interface
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/DLTI/DLTI.cpp

 
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecAttr
 //===----------------------------------------------------------------------===//
 //
 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
 constexpr const StringLiteral
     mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
+constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutStackAlignmentKey;
 
 namespace mlir {
 namespace impl {
 class DataLayoutSpecStorage : public AttributeStorage {
 public:
   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
 
   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)

 }
 
 StringAttr
 DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
   return Builder(context).getStringAttr(
       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
 ///               | attr `,` attr-list
 DataLayoutSpecAttr DataLayoutSpecAttr::parse(AsmParser &parser) {
   if (failed(parser.parseLess()))
     return {};
 
if (entryName == DLTIDialect::kDataLayoutEndiannessKey) {
           (value.getValue() == DLTIDialect::kDataLayoutEndiannessBig ||
            value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle))
         return success();
       return emitError(loc) << "'" << entryName
                             << "' data layout entry is expected to be either '"
                             << 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;
   }
 };
 } // namespace
 
 void DLTIDialect::initialize() {
   addAttributes<DataLayoutEntryAttr, DataLayoutSpecAttr>();

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

 mlir::detail::getDefaultAllocaMemorySpace(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.
+unsigned
+mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface())
+    return 0;
+
+  auto attr = entry.getValue().cast<IntegerAttr>();
+  return attr.getValue().getZExtValue();
+}
+
 DataLayoutEntryList
 mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
                                    TypeID typeID) {
   return llvm::to_vector<4>(llvm::make_filter_range(
       entries, [typeID](DataLayoutEntryInterface entry) {
         auto type = entry.getKey().dyn_cast<Type>();
         return type && type.getTypeID() == typeID;
       }));
assert((!op || !op.getDataLayoutSpec()) &&
            "combine attributes?)");
   }
 }
 
 mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
 
 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);
 #endif
 }
 
 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);
 #endif
 }
 
 mlir::DataLayout mlir::DataLayout::closest(Operation *op) {
   // Search the closest parent either being a module operation or implementing
entry = originalLayout.getSpecForIdentifier(
         originalLayout.getAllocaMemorySpaceIdentifier(context));
   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
     allocaMemorySpace = iface.getAllocaMemorySpace(entry);
   else
     allocaMemorySpace = detail::getDefaultAllocaMemorySpace(entry);
   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<DataLayoutOpInterface>(scope))
+    stackAlignment = iface.getStackAlignment(entry);
+  else
+    stackAlignment = detail::getDefaultStackAlignment(entry);
+  return *stackAlignment;
+}
+
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecInterface
 //===----------------------------------------------------------------------===//
 
 void DataLayoutSpecInterface::bucketEntriesByType(
     DenseMap<TypeID, DataLayoutEntryList> &types,
     DenseMap<StringAttr, DataLayoutEntryInterface> &ids) {
   for (DataLayoutEntryInterface entry : getEntries()) {

mlir/lib/Target/LLVMIR/DataLayoutImporter.h

 
 namespace detail {
 
 /// Returns a supported MLIR floating point type of the given bit width or
 /// null if the bit width is not supported.
 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,
                      const llvm::DataLayout &llvmDataLayout)
       : context(context) {
     translateDataLayout(llvmDataLayout);
   }
 
private:
 
   /// Adds an endianness entry if there is none yet.
   LogicalResult tryToEmplaceEndiannessEntry(StringRef endianness,
                                             StringRef token);
 
   /// 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<StringRef> unhandledTokens;
   DenseMap<StringAttr, DataLayoutEntryInterface> keyEntries;
   DenseMap<TypeAttr, DataLayoutEntryInterface> typeEntries;
   MLIRContext *context;
   DataLayoutSpecInterface dataLayout;
 };
 
 } // namespace detail
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_LIB_TARGET_LLVMIR_DATALAYOUTIMPORTER_H_

mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp

DataLayoutImporter::tryToEmplaceAllocaAddrSpaceEntry(StringRef token) {
   if (*space == 0)
     return success();
   OpBuilder builder(context);
   keyEntries.try_emplace(
       key, DataLayoutEntryAttr::get(key, builder.getUI32IntegerAttr(*space)));
   return success();
 }
 
+LogicalResult
+DataLayoutImporter::tryToEmplaceStackAlignmentEntry(StringRef token) {
+  auto key =
+      StringAttr::get(context, DLTIDialect::kDataLayoutStackAlignmentKey);
+  if (keyEntries.count(key))
+    return success();
+
+  FailureOr<unsigned> 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 = {};
 
   // Transform the data layout to its string representation and append the
   // default data layout string specified in the language reference
   // (https://llvm.org/docs/LangRef.html#data-layout). The translation then
   // parses the string and ignores the default value if a specific kind occurs
if (*prefix == "E") {
       continue;
     }
     // Parse the alloca address space.
     if (*prefix == "A") {
       if (failed(tryToEmplaceAllocaAddrSpaceEntry(token)))
         return;
       continue;
     }
+    // Parse the stack alignment.
+    if (*prefix == "S") {
+      if (failed(tryToEmplaceStackAlignmentEntry(token)))
+        return;
+      continue;
+    }
     // Parse integer alignment specifications.
     if (*prefix == "i") {
       FailureOr<unsigned> width = tryToParseInt(token);
       if (failed(width))
         return;
 
       Type type = IntegerType::get(context, *width);
       if (failed(tryToEmplaceAlignmentEntry(type, token)))

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

translateDataLayout(DataLayoutSpecInterface attribute,
   for (DataLayoutEntryInterface entry : attribute.getEntries()) {
     auto key = entry.getKey().dyn_cast<StringAttr>();
     if (!key)
       continue;
     if (key.getValue() == DLTIDialect::kDataLayoutEndiannessKey) {
       auto value = entry.getValue().cast<StringAttr>();
       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<IntegerAttr>();
-      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;
+      continue;
+    if (key.getValue() == DLTIDialect::kDataLayoutStackAlignmentKey) {
+      auto value = entry.getValue().cast<IntegerAttr>();
+      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;
     return failure();
   }
 
   // Go through the list of entries to check which types are explicitly
   // specified in entries. Where possible, data layout queries are used instead

mlir/test/Dialect/LLVMIR/layout.mlir

 // RUN: mlir-opt --test-data-layout-query --split-input-file --verify-diagnostics %s | FileCheck %s
 
 module {
   // CHECK: @no_spec
   func.func @no_spec() {
     // 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
     // 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
   }
 }
 
 // -----
 
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.ptr, dense<[32, 32, 64]> : vector<3xi32>>,
   #dlti.dl_entry<!llvm.ptr<5>, dense<[64, 64, 64]> : vector<3xi32>>,
   #dlti.dl_entry<!llvm.ptr<4>, 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() {
     // 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
     // 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
+    // CHECK: stack_alignment = 128
   "test.data_layout_query"() : () -> !llvm.ptr<4>
     return
   }
 }
 
 // -----
 
 // expected-error@below {{expected layout attribute for '!llvm.ptr' to be a dense integer elements attribute with 3 or 4 elements}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<

mlir/test/Target/LLVMIR/Import/data-layout.ll

 ; CHECK: #dlti.dl_spec<
 ; CHECK-DAG:   #dlti.dl_entry<"dlti.endianness", "little">
 ; CHECK-DAG:   #dlti.dl_entry<i64, dense<64> : vector<2xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<f80, dense<128> : vector<2xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<i8, dense<8> : vector<2xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<!llvm.ptr<270>, dense<[32, 64, 64, 32]> : vector<4xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<!llvm.ptr<272>, 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"
 
 ; // -----
 
 ; CHECK: dlti.dl_spec =
 ; CHECK: #dlti.dl_spec<
 ; CHECK-DAG:   #dlti.dl_entry<"dlti.endianness", "big">
 ; CHECK-DAG:   #dlti.dl_entry<!llvm.ptr<270>, dense<[16, 32, 64, 128]> : vector<4xi32>>
+; CHECK-DAG:   #dlti.dl_entry<!llvm.ptr<271>, dense<[16, 32, 64, 16]> : vector<4xi32>>
 ; CHECK-DAG:   #dlti.dl_entry<"dlti.alloca_memory_space", 1 : ui32>
 ; CHECK-DAG:   #dlti.dl_entry<i64, dense<[64, 128]> : 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"

mlir/test/Target/LLVMIR/data-layout.mlir

 // RUN: mlir-translate -mlir-to-llvmir %s -split-input-file -verify-diagnostics | FileCheck %s
 
 // 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<index, 64>,
 #dlti.dl_entry<i64, dense<[64,128]> : vector<2xi32>>,
 #dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>,
 #dlti.dl_entry<!llvm.ptr, dense<[32,64,128]> : vector<3xi32>>,
 #dlti.dl_entry<!llvm.ptr<1>, dense<[32,32,32,64]> : vector<4xi32>>
 >} {
   llvm.func @foo() {
     llvm.return
   }
 }
 
 // -----
 
+// 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<ui64, dense<[64,128]> : vector<2xi32>>>
 } {}
 
 // -----
 
 // expected-error@below {{unsupported type in data layout: 'bf16'}}

mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp

func.walk([&](test::DataLayoutQueryOp op) {
         return;
 
       const DataLayout &layout = layouts.getAbove(op);
       unsigned size = layout.getTypeSize(op.getType());
       unsigned bitsize = layout.getTypeSizeInBits(op.getType());
       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)),
            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)),
            builder.getNamedAttr("alloca_memory_space",
                                 allocaMemorySpace == Attribute()
                                     ? builder.getUI32IntegerAttr(0)
-                                    : allocaMemorySpace)
-
-          });
+                                    : allocaMemorySpace),
+           builder.getNamedAttr("stack_alignment",
+                                builder.getIndexAttr(stackAlignment))});
     });
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
 } // namespace test
 } // namespace mlir

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

 #include <gtest/gtest.h>
 
 using namespace mlir;
 
 namespace {
 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.
 class DataLayoutSpecStorage : public AttributeStorage {
 public:
   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
 
   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
struct CustomDataLayoutSpec
   combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
     return *this;
   }
   DataLayoutEntryListRef getEntries() const { return getImpl()->entries; }
   LogicalResult verifySpec(Location loc) { return success(); }
   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
 /// than once. Handy to check if the cache works.
 struct SingleQueryType
     : public Type::TypeBase<SingleQueryType, Type, TypeStorage,
                             DataLayoutTypeInterface::Trait> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SingleQueryType)
module {}
   EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u);
   EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 16u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
   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());
+  EXPECT_EQ(layout.getStackAlignment(), 0u);
 }
 
 TEST(DataLayout, NullSpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() : () -> ()
   )MLIR";
 
   DialectRegistry registry;
"dltest.op_with_layout"() : () -> ()
   EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
   EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
   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());
+  EXPECT_EQ(layout.getStackAlignment(), 0u);
 }
 
 TEST(DataLayout, EmptySpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
   )MLIR";
 
   DialectRegistry registry;
"dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
   EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
   EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
   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());
+  EXPECT_EQ(layout.getStackAlignment(), 0u);
 }
 
 TEST(DataLayout, SpecWithEntries) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<
   #dlti.dl_entry<i42, 5>,
   #dlti.dl_entry<i16, 6>,
-  #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";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);
 
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
> } : () -> ()
   EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 256u);
   EXPECT_EQ(layout.getTypeSizeInBits(Float32Type::get(&ctx)), 256u);
   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u);
   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));
+  EXPECT_EQ(layout.getStackAlignment(), 128u);
 }
 
 TEST(DataLayout, Caching) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
   )MLIR";
 
   DialectRegistry registry;