[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

diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 77a0367..cf78b23 100644 (file)
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -101,6 +101,9 @@ public:
   /// 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
index a1f9eb3..37187bf 100644 (file)
--- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
+++ b/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;

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 3d4d379..e6e31d1 100644 (file)
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/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 @@ public:
   /// 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 @@ private:
 
   /// Cache for alloca memory space.
   mutable std::optional<Attribute> allocaMemorySpace;
+
+  /// Cache for stack alignment.
+  mutable std::optional<unsigned> stackAlignment;
 };
 
 } // namespace mlir

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 2d9d6ac..bac8c23 100644 (file)
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/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); }];

diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 86cc3ea..889778a 100644 (file)
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/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 @@ public:
                             << 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
index ba6a321..ae8e680 100644 (file)
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/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
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
index af480c5..d195906 100644 (file)
--- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
+++ b/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);

diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
index 2b2fa66..5c16062 100644 (file)
--- a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
+++ b/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 @@ private:
   /// 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;

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 7159aa4..a08bb37 100644 (file)
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/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;

diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index ba388c1..363222b 100644 (file)
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/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
   }
 }

diff --git a/mlir/test/Target/LLVMIR/Import/data-layout.ll b/mlir/test/Target/LLVMIR/Import/data-layout.ll
index e29e174..b46f5fc 100644 (file)
--- a/mlir/test/Target/LLVMIR/Import/data-layout.ll
+++ b/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"

diff --git a/mlir/test/Target/LLVMIR/data-layout.mlir b/mlir/test/Target/LLVMIR/data-layout.mlir
index 5840bcf..74941d8 100644 (file)
--- 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<index, 64>,
 #dlti.dl_entry<i64, dense<[64,128]> : vector<2xi32>>,
 #dlti.dl_entry<f80, dense<[128,256]> : vector<2xi32>>,
@@ -21,6 +25,22 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
 
 // -----
 
+// 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>>>

diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index 9e71bd4..c326a74 100644 (file)
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -41,6 +41,7 @@ struct TestDataLayoutQuery
       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 @@ struct TestDataLayoutQuery
            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
index c35e176..7e0a8f5 100644 (file)
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -24,6 +24,8 @@ 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.
@@ -67,6 +69,9 @@ struct CustomDataLayoutSpec
   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 @@ module {}
   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 @@ TEST(DataLayout, NullSpec) {
   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 @@ TEST(DataLayout, EmptySpec) {
   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 @@ TEST(DataLayout, SpecWithEntries) {
 "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";
 
@@ -365,6 +374,7 @@ TEST(DataLayout, SpecWithEntries) {
   EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
 
   EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5));
+  EXPECT_EQ(layout.getStackAlignment(), 128u);
 }
 
 TEST(DataLayout, Caching) {