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

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.

Reviewed By: Dinistro

Differential Revision: https://reviews.llvm.org/D147332

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

@@ -101,6 +101,9 @@ class DataLayoutSpecAttr
   /// 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);
 

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

@@ -39,6 +39,9 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
   }];
 
   let useDefaultAttributePrinterParser = 1;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -60,6 +60,10 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
 /// 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 @@ class DataLayout {
   /// 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 @@ class DataLayout {
 
   /// Cache for alloca memory space.
   mutable std::optional<Attribute> allocaMemorySpace;
+
+  /// Cache for stack alignment.
+  mutable std::optional<unsigned> stackAlignment;
 };
 
 } // namespace mlir

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -112,6 +112,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> {
       /*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 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         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); }];

mlir/lib/Dialect/DLTI/DLTI.cpp

@@ -108,6 +108,7 @@ void DataLayoutEntryAttr::print(AsmPrinter &os) const {
 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 @@ DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
       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 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
                             << 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;
   }

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -225,6 +225,17 @@ mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
   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 value = entry.getValue().cast<IntegerAttr>();
+  return value.getValue().getZExtValue();
+}
+
 DataLayoutEntryList
 mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
                                    TypeID typeID) {
@@ -359,7 +370,7 @@ 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);
@@ -368,7 +379,7 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
 
 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 @@ mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
   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
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp

@@ -183,6 +183,26 @@ DataLayoutImporter::tryToEmplaceAllocaAddrSpaceEntry(StringRef token) {
   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 = {};
@@ -230,6 +250,12 @@ void DataLayoutImporter::translateDataLayout(
         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);

mlir/lib/Target/LLVMIR/DataLayoutImporter.h

@@ -38,11 +38,11 @@ namespace detail {
 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 @@ class DataLayoutImporter {
   /// 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;

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -75,17 +75,28 @@ translateDataLayout(DataLayoutSpecInterface attribute,
       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;
+      layoutStream.flush();
+      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;

mlir/test/Dialect/LLVMIR/layout.mlir

@@ -8,18 +8,21 @@ module {
     // 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 @@ 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() {
@@ -40,31 +44,36 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // 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
   }
 }

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

@@ -28,6 +28,7 @@ target datalayout = ""
 ; 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"
 
 ; // -----
@@ -36,13 +37,15 @@ target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16
 ; 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"