[Mlir-commits] [mlir] 382eb7c - [mlir] Add alloca address space handling to the data layout subsystem

Jan Sjodin llvmlistbot at llvm.org
Tue Mar 21 06:50:42 PDT 2023


Author: Jan Sjodin
Date: 2023-03-21T09:37:08-04:00
New Revision: 382eb7c2c71648fe15fe568ba67f93d3b53beb08

URL: https://github.com/llvm/llvm-project/commit/382eb7c2c71648fe15fe568ba67f93d3b53beb08
DIFF: https://github.com/llvm/llvm-project/commit/382eb7c2c71648fe15fe568ba67f93d3b53beb08.diff

LOG: [mlir] Add alloca address space handling to the data layout subsystem

This patch adds alloca address space information to the data layout interface
and implementation in the DLTI dialect. This is needed for targets that use
separate address spaces for local/stack data.

Reviewed By: ftynse, krzysz00

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/DLTI/DLTI.h
    mlir/include/mlir/Dialect/DLTI/DLTIBase.td
    mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
    mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
    mlir/lib/Dialect/DLTI/DLTI.cpp
    mlir/lib/Interfaces/DataLayoutInterfaces.cpp
    mlir/lib/Target/LLVMIR/ModuleImport.cpp
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    mlir/test/Dialect/LLVMIR/layout.mlir
    mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
    mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 1aa9ab5e250a..77a036754f96 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -98,6 +98,9 @@ class DataLayoutSpecAttr
   /// Returns the list of entries.
   DataLayoutEntryListRef getEntries() const;
 
+  /// Returns the alloca memory space identifier.
+  StringAttr getAllocaMemorySpaceIdentifier(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 866cc934e05c..a1f9eb361d71 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
@@ -23,9 +23,11 @@ def DLTI_Dialect : Dialect {
   }];
 
   let extraClassDeclaration = [{
+    // Top level attribute name.
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAttrName = "dlti.dl_spec";
 
+    // Constants used in entries.
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessKey = "dlti.endianness";
 
@@ -34,6 +36,9 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessLittle = "little";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
   }];
 
   let useDefaultAttributePrinterParser = 1;

diff  --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 5cc80a06f363..3d4d3796ecca 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -56,6 +56,10 @@ 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);
+
 /// 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.
@@ -159,6 +163,9 @@ class DataLayout {
   /// 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;
+
 private:
   /// Combined layout spec at the given scope.
   const DataLayoutSpecInterface originalLayout;
@@ -180,6 +187,9 @@ class DataLayout {
   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;
 };
 
 } // namespace mlir

diff  --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 0ca0f9e22191..2d9d6acbedae 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -106,6 +106,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> {
       /*methodName=*/"getEntries",
       /*args=*/(ins)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the alloca memory space identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getAllocaMemorySpaceIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     // Implementations may override this if they have an efficient lookup
     // mechanism.
     InterfaceMethod<
@@ -256,6 +262,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
                                                             params);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the memory space used by the ABI computed "
+                      "using the relevant entries. The data layout object "
+                      "can be used for recursive queries.",
+      /*retTy=*/"::mlir::Attribute",
+      /*methodName=*/"getAllocaMemorySpace",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultAllocaMemorySpace(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 eaf6f1e619a0..86cc3ea76603 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -106,6 +106,8 @@ void DataLayoutEntryAttr::print(AsmPrinter &os) const {
 //===----------------------------------------------------------------------===//
 //
 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
+constexpr const StringLiteral
+    mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
 
 namespace mlir {
 namespace impl {
@@ -273,6 +275,12 @@ DataLayoutEntryListRef DataLayoutSpecAttr::getEntries() const {
   return getImpl()->entries;
 }
 
+StringAttr
+DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutAllocaMemorySpaceKey);
+}
+
 /// Parses an attribute with syntax
 ///   attr ::= `#target.` `dl_spec` `<` attr-list? `>`
 ///   attr-list ::= attr
@@ -329,6 +337,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
                             << DLTIDialect::kDataLayoutEndiannessBig << "' or '"
                             << DLTIDialect::kDataLayoutEndiannessLittle << "'";
     }
+    if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey)
+      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 141ec9e44166..ba6a321360fa 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -213,6 +213,18 @@ unsigned mlir::detail::getDefaultPreferredAlignment(
   reportMissingDataLayout(type);
 }
 
+// Returns the memory space used for allocal operations if specified in the
+// given entry. If the entry is empty the default memory space represented by
+// an empty attribute is returned.
+Attribute
+mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface()) {
+    return Attribute();
+  }
+
+  return entry.getValue();
+}
+
 DataLayoutEntryList
 mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
                                    TypeID typeID) {
@@ -346,7 +358,8 @@ void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
 mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
 
 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
-    : originalLayout(getCombinedDataLayout(op)), scope(op) {
+    : originalLayout(getCombinedDataLayout(op)), scope(op),
+      allocaMemorySpace(std::nullopt) {
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
@@ -354,7 +367,8 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
 }
 
 mlir::DataLayout::DataLayout(ModuleOp op)
-    : originalLayout(getCombinedDataLayout(op)), scope(op) {
+    : originalLayout(getCombinedDataLayout(op)), scope(op),
+      allocaMemorySpace(std::nullopt) {
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
@@ -456,6 +470,22 @@ unsigned mlir::DataLayout::getTypePreferredAlignment(Type t) const {
   });
 }
 
+mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
+  checkValid();
+  MLIRContext *context = scope->getContext();
+  if (allocaMemorySpace)
+    return *allocaMemorySpace;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    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;
+}
+
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecInterface
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index eecb85527513..a45768726bc4 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -20,6 +20,7 @@
 
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
@@ -179,9 +180,10 @@ mlir::translateDataLayout(const llvm::DataLayout &dataLayout,
   // Remaining unhandled default layout defaults
   // e (little endian if not set)
   // p[n]:64:64:64 (non zero address spaces have 64-bit properties)
+  // Alloca address space defaults to 0.
   std::string append =
       "p:64:64:64-S0-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f16:16:16-f64:"
-      "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64";
+      "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64-A0";
   if (layoutstr.empty())
     layoutstr = append;
   else
@@ -227,6 +229,18 @@ mlir::translateDataLayout(const llvm::DataLayout &dataLayout,
           StringAttr::get(context, DLTIDialect::kDataLayoutEndiannessKey),
           value);
       entries.emplace_back(entry);
+    } else if (symbol == 'A') {
+      unsigned addressSpace;
+      if (parameter.getAsInteger(/*Radix=*/10, addressSpace))
+        return nullptr;
+      // Skip storing if generic address space is defined.
+      if (addressSpace != 0) {
+        auto entry = DataLayoutEntryAttr::get(
+            StringAttr::get(context,
+                            DLTIDialect::kDataLayoutAllocaMemorySpaceKey),
+            mlir::Builder(context).getUI32IntegerAttr(addressSpace));
+        entries.emplace_back(entry);
+      }
     }
   }
 

diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 9b0bfa22c73a..ca5fb7dda384 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -79,6 +79,15 @@ translateDataLayout(DataLayoutSpecInterface attribute,
       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();
+      }
+      continue;
+    }
     emitError(*loc) << "unsupported data layout key " << key;
     return failure();
   }

diff  --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
index c2f162dd6175..e5c8c0bd86db 100644
--- a/mlir/test/Dialect/LLVMIR/layout.mlir
+++ b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -4,37 +4,44 @@ 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
     "test.data_layout_query"() : () -> !llvm.ptr<i8>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<i32>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<bf16>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 0
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
     // CHECK: alignment = 8
-	// CHECK: bitsize = 64
+    // CHECK: alloca_memory_space = 0
+    // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<5>
@@ -47,47 +54,56 @@ module {
 module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<!llvm.ptr<i8>, dense<[32, 32, 64]> : vector<3xi32>>,
   #dlti.dl_entry<!llvm.ptr<i8, 5>, dense<[64, 64, 64]> : vector<3xi32>>,
-  #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 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>
 >} {
   // CHECK: @spec
   func.func @spec() {
     // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<i8>
     // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<i32>
     // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<bf16>
     // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
     // CHECK: alignment = 4
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
     // CHECK: alignment = 8
+    // CHECK: alloca_memory_space = 5
     // CHECK: bitsize = 64
     // CHECK: preferred = 8
     // CHECK: size = 8
     "test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
     // CHECK: alignment = 4
-	  // CHECK: bitsize = 32
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
     "test.data_layout_query"() : () -> !llvm.ptr<3>
     // CHECK: alignment = 8
-	  // CHECK: bitsize = 32
+    // CHECK: alloca_memory_space = 5
+    // CHECK: bitsize = 32
     // CHECK: preferred = 8
     // CHECK: size = 4
 	"test.data_layout_query"() : () -> !llvm.ptr<4>

diff  --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index a7fde171dd85..9e71bd4e999f 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -40,11 +40,18 @@ struct TestDataLayoutQuery
       unsigned bitsize = layout.getTypeSizeInBits(op.getType());
       unsigned alignment = layout.getTypeABIAlignment(op.getType());
       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
+      Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
       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("preferred", builder.getIndexAttr(preferred)),
+           builder.getNamedAttr("alloca_memory_space",
+                                allocaMemorySpace == Attribute()
+                                    ? builder.getUI32IntegerAttr(0)
+                                    : allocaMemorySpace)
+
+          });
     });
   }
 };

diff  --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index 31176aa07d2b..c35e17603522 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -22,6 +22,8 @@ using namespace mlir;
 
 namespace {
 constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
+constexpr static llvm::StringLiteral kAllocaKeyName =
+    "dltest.alloca_memory_space";
 
 /// Trivial array storage for the custom data layout spec attribute, just a list
 /// of entries.
@@ -62,6 +64,9 @@ struct CustomDataLayoutSpec
   }
   DataLayoutEntryListRef getEntries() const { return getImpl()->entries; }
   LogicalResult verifySpec(Location loc) { return success(); }
+  StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
+    return Builder(context).getStringAttr(kAllocaKeyName);
+  }
 };
 
 /// A type subject to data layout that exits the program if it is queried more
@@ -104,6 +109,15 @@ struct SingleQueryType
     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.
@@ -260,6 +274,8 @@ module {}
   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());
 }
 
 TEST(DataLayout, NullSpec) {
@@ -275,6 +291,7 @@ TEST(DataLayout, NullSpec) {
   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 +300,8 @@ TEST(DataLayout, NullSpec) {
   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());
 }
 
 TEST(DataLayout, EmptySpec) {
@@ -306,13 +325,16 @@ TEST(DataLayout, EmptySpec) {
   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());
 }
 
 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>
 > } : () -> ()
   )MLIR";
 
@@ -341,6 +363,8 @@ TEST(DataLayout, SpecWithEntries) {
   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));
 }
 
 TEST(DataLayout, Caching) {


        


More information about the Mlir-commits mailing list