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

[jsjodin]

Summary

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.

[krzysz00]

Since we want to be able to use this whenever someone's creating a memref::AllocaOp, we want a generic Attribute-valued memory space key.

However, we could also have both the unsigned-valued allocaAddressSpace and the Attribute-valued allocaMemorySpace as separate keys, with addressSpace used in LLVM lowering on down and allocaMemorySpace used in MLIR proper. Though that'd require some documentation to make sure people understand which is which ...

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

102

1. For consistency with memref, can we call this getAllocaMemorySpaceIdentifier()

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

61

I'd argue this should also be Attribute, and that, when dealing with the *ToLLVMs, it should be sent through the attribute conversion process.

192

mutable Attribute memorySpace

mlir/lib/Dialect/DLTI/DLTI.cpp

340

Strong disagree with the restriction to integer attributes, consider that this needs to work for #gpu.address_space<Private>

[krzysz00]

... in further comments, if we want to go with the address space and memory space as distinct data layout entries system, I'd be willing to review this patch on its own and leave the memory space part to another patch.

[jsjodin]

In D144657#4148055, @krzysz00 wrote:

... in further comments, if we want to go with the address space and memory space as distinct data layout entries system, I'd be willing to review this patch on its own and leave the memory space part to another patch.

It would probably be better to have a generic entry system that can work for anything

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

61

Yes, that is more general. I will work on an updated patch.

192

Yes this makes sense. Each dialect (like LLVMIR) that uses the data layout could have their own convenience functions to handle the values. Where would the best way to put these functions?

[jsjodin]

Also, how would the default case be handled (currently returns 0)?

[krzysz00]

Re the default value of getAllocaMemorySpcae(), that would be the null attribute, aka {} aka nullptr, which is the current default.

[jsjodin]

Use Attribute instead of integers to encode memory spaces. Change address space -> memory space.

[krzysz00]

Looks good to me - though I'm not sure if we want @ftynse or someone else more familiar with DTLI to take a look before landing.

[jsjodin]

In D144657#4158826, @krzysz00 wrote:

Looks good to me - though I'm not sure if we want @ftynse or someone else more familiar with DTLI to take a look before landing.

Yes, I will let @ftynse approve.

[ftynse]

mlir/lib/Target/LLVMIR/ModuleImport.cpp

236

Ultra-nit: terminate comments with a dot.

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

278

Can we have a test with non-null memory space?

[ftynse]

Sorry for delay, already-approved patches don't show up in the review stream :(

[jsjodin]

In D144657#4206620, @ftynse wrote:

Sorry for delay, already-approved patches don't show up in the review stream :(

No problem, I suspected that might be the case.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

236

Ultra-nit: terminate comments with a dot.

Will fix.

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

278

Sure, I will work on another test case.

[jsjodin]

Added unittest with non-null memory space. Fixed comments to end with '.'.

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

public:
   /// from outermost to innermost. This overwrites the older entries with the
   /// same key as the newer entries if the entries are compatible. Returns null
   /// if the specifications are not compatible.
   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;
+
   /// 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

def DLTI_Dialect : Dialect {
 
   let description = [{
     The Data Layout and Target Information (DLTI) dialect is intended to hold
     attributes and other components pertaining to descriptions of in-memory
     data layout and compilation targets.
   }];
 
   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";
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessBig = "big";
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessLittle = "little";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
   }];
 
   let useDefaultAttributePrinterParser = 1;
 }
 
 def DLTI_DataLayoutEntryAttr : DialectAttr<
     DLTI_Dialect,
     CPred<"$_self.isa<::mlir::DataLayoutEntryAttr>()">,

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

 
 /// Default handler for the preferred alignemnt request. Computes results for
 /// built-in types and dispatches to the DataLayoutTypeInterface for other
 /// types.
 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.
 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 getTypeSizeInBits(Type t) const;
 
   /// Returns the required alignment of the given type in the current scope.
   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;
+
 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;
 };
 
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

InterfaceMethod<
       /*args=*/(ins "::llvm::ArrayRef<::mlir::DataLayoutSpecInterface>":$specs)
     >,
     InterfaceMethod<
       /*description=*/"Returns the list of layout entries.",
       /*retTy=*/"::mlir::DataLayoutEntryListRef",
       /*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<
       /*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<
                     "const ::mlir::DataLayout &":$dataLayout,
                     "::mlir::DataLayoutEntryListRef":$params),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return ::mlir::detail::getDefaultPreferredAlignment(type, dataLayout,
                                                             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); }];
 }
 
 //===----------------------------------------------------------------------===//
 // Type interface
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/DLTI/DLTI.cpp

void DataLayoutEntryAttr::print(AsmPrinter &os) const {
   os << ", " << getValue() << ">";
 }
 
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecAttr
 //===----------------------------------------------------------------------===//
 //
 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
+constexpr const StringLiteral
+    mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
 
 namespace mlir {
 namespace impl {
 class DataLayoutSpecStorage : public AttributeStorage {
 public:
   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
 
   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
DataLayoutSpecAttr::combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
 
   return DataLayoutSpecAttr::get(getContext(), entries);
 }
 
 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
 ///               | 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)
+      return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;
   }
 };
 } // namespace
 
 void DLTIDialect::initialize() {
   addAttributes<DataLayoutEntryAttr, DataLayoutSpecAttr>();
   addInterfaces<TargetDataLayoutInterface>();

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

return getDefaultPreferredAlignment(ctype.getElementType(), dataLayout,
                                         params);
 
   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
     return typeInterface.getPreferredAlignment(dataLayout, params);
 
   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) {
   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()) &&
            "could not compute layout information for an op (failed to "
            "combine attributes?)");
   }
 }
 
 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);
 #endif
 }
 
 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);
 #endif
 }
 
 mlir::DataLayout mlir::DataLayout::closest(Operation *op) {
   // Search the closest parent either being a module operation or implementing
return cachedLookup(t, preferredAlignments, [&](Type ty) {
     if (originalLayout)
       list = originalLayout.getSpecForType(ty.getTypeID());
     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
       return iface.getTypePreferredAlignment(ty, *this, list);
     return detail::getDefaultPreferredAlignment(ty, *this, list);
   });
 }
 
+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
 //===----------------------------------------------------------------------===//
 
 void DataLayoutSpecInterface::bucketEntriesByType(
     DenseMap<TypeID, DataLayoutEntryList> &types,
     DenseMap<StringAttr, DataLayoutEntryInterface> &ids) {
   for (DataLayoutEntryInterface entry : getEntries()) {

mlir/lib/Target/LLVMIR/ModuleImport.cpp

 #include "mlir/Target/LLVMIR/Import.h"
 
 #include "AttrKindDetail.h"
 #include "DebugImporter.h"
 #include "LoopAnnotationImporter.h"
 
 #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"
 
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/IR/Constants.h"

 mlir::translateDataLayout(const llvm::DataLayout &dataLayout,
                           MLIRContext *context) {
   assert(context && "expected MLIR context");
   std::string layoutstr = dataLayout.getStringRepresentation();
 
   // 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
     layoutstr = layoutstr + "-" + append;
 
   StringRef layout(layoutstr);
 
   SmallVector<DataLayoutEntryInterface> entries;
while (!layout.empty()) {
     } else if (symbol == 'e' || symbol == 'E') {
       auto value = StringAttr::get(
           context, symbol == 'e' ? DLTIDialect::kDataLayoutEndiannessLittle
                                  : DLTIDialect::kDataLayoutEndiannessBig);
       auto entry = DataLayoutEntryAttr::get(
           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);
+      }
     }
   }
 
   return DataLayoutSpecAttr::get(context, entries);
 }
 
 /// Get a topologically sorted list of blocks for the given function.
 static SetVector<llvm::BasicBlock *>

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (DataLayoutEntryInterface entry : attribute.getEntries()) {
     if (key.getValue() == DLTIDialect::kDataLayoutEndiannessKey) {
       auto value = entry.getValue().cast<StringAttr>();
       bool isLittleEndian =
           value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle;
       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();
+      }
+      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
   // of directly inspecting the entries.
   for (DataLayoutEntryInterface entry : attribute.getEntries()) {

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
     "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>
     return
   }
 }
 
 // -----
 
 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>
     return
   }
 }
 
 // -----

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

func.walk([&](test::DataLayoutQueryOp op) {
       if (op->getAttr("size"))
         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();
       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)
+
+          });
     });
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
 } // namespace test
 } // namespace mlir

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

 #include "mlir/Parser/Parser.h"
 
 #include <gtest/gtest.h>
 
 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.
 class DataLayoutSpecStorage : public AttributeStorage {
 public:
   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
 
   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
static CustomDataLayoutSpec get(MLIRContext *ctx,
     return Base::get(ctx, entries);
   }
   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);
+  }
 };
 
 /// 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)
unsigned getPreferredAlignment(const DataLayout &layout,
                                  DataLayoutEntryListRef params) {
     static bool executed = false;
     if (executed)
       llvm::report_fatal_error("repeated call");
 
     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.
 struct TypeNoLayout : public Type::TypeBase<TypeNoLayout, Type, TypeStorage> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeNoLayout)
 
   using Base::Base;
 
module {}
   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u);
   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());
 }
 
 TEST(DataLayout, NullSpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() : () -> ()
   )MLIR";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);
 
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
   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);
   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());
 }
 
 TEST(DataLayout, EmptySpec) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
   )MLIR";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);
 
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
   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);
   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());
 }
 
 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";
 
   DialectRegistry registry;
   registry.insert<DLTIDialect, DLTestDialect>();
   MLIRContext ctx(registry);
 
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
> } : () -> ()
   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u);
   EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u);
   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));
 }
 
 TEST(DataLayout, Caching) {
   const char *ir = R"MLIR(
 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
   )MLIR";
 
   DialectRegistry registry;