--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp	Wed Nov 09 17:47:54 2022 +0900
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp	Fri Aug 18 09:04:13 2023 +0900
@@ -22,6 +22,10 @@
 
 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.
@@ -62,6 +66,12 @@
   }
   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
@@ -104,6 +114,15 @@
     executed = true;
     return 4;
   }
+
+  Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) {
+    static bool executed = false;
+    if (executed)
+      llvm::report_fatal_error("repeated call");
+
+    executed = true;
+    return Attribute();
+  }
 };
 
 /// A types that is not subject to data layout.
@@ -133,16 +152,16 @@
   static unsigned getTypeSizeInBits(Type type, const DataLayout &dataLayout,
                                     DataLayoutEntryListRef params) {
     // Make a recursive query.
-    if (type.isa<FloatType>())
+    if (isa<FloatType>(type))
       return dataLayout.getTypeSizeInBits(
           IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth()));
 
     // Handle built-in types that are not handled by the default process.
-    if (auto iType = type.dyn_cast<IntegerType>()) {
+    if (auto iType = dyn_cast<IntegerType>(type)) {
       for (DataLayoutEntryInterface entry : params)
-        if (entry.getKey().dyn_cast<Type>() == type)
+        if (llvm::dyn_cast_if_present<Type>(entry.getKey()) == type)
           return 8 *
-                 entry.getValue().cast<IntegerAttr>().getValue().getZExtValue();
+                 cast<IntegerAttr>(entry.getValue()).getValue().getZExtValue();
       return 8 * iType.getIntOrFloatBitWidth();
     }
 
@@ -198,7 +217,7 @@
   void printAttribute(Attribute attr,
                       DialectAsmPrinter &printer) const override {
     printer << "spec<";
-    llvm::interleaveComma(attr.cast<CustomDataLayoutSpec>().getEntries(),
+    llvm::interleaveComma(cast<CustomDataLayoutSpec>(attr).getEntries(),
                           printer);
     printer << ">";
   }
@@ -225,7 +244,7 @@
   }
 
   void printType(Type type, DialectAsmPrinter &printer) const override {
-    if (type.isa<SingleQueryType>())
+    if (isa<SingleQueryType>(type))
       printer << "single_query";
     else
       printer << "no_layout";
@@ -260,6 +279,9 @@
   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) {
@@ -275,6 +297,7 @@
   auto op =
       cast<DataLayoutOpInterface>(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 +306,9 @@
   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) {
@@ -306,13 +332,18 @@
   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<i16, 6>,
+  #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>,
+  #dlti.dl_entry<"dltest.stack_alignment", 128 : i32>
 > } : () -> ()
   )MLIR";
 
@@ -341,6 +372,9 @@
   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) {