Index: mlir/include/mlir/Dialect/DLTI/DLTI.h
===================================================================
--- mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -98,6 +98,9 @@
   /// 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);
 
Index: mlir/include/mlir/Dialect/DLTI/DLTIBase.td
===================================================================
--- mlir/include/mlir/Dialect/DLTI/DLTIBase.td
+++ mlir/include/mlir/Dialect/DLTI/DLTIBase.td
@@ -23,9 +23,11 @@
   }];
 
   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 @@
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessLittle = "little";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
   }];
 
   let useDefaultAttributePrinterParser = 1;
Index: mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
===================================================================
--- mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -56,6 +56,10 @@
 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 @@
   /// 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 @@
   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
Index: mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
===================================================================
--- mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -106,6 +106,12 @@
       /*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 @@
                                                             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); }];
Index: mlir/lib/Dialect/DLTI/DLTI.cpp
===================================================================
--- mlir/lib/Dialect/DLTI/DLTI.cpp
+++ mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -106,6 +106,8 @@
 //===----------------------------------------------------------------------===//
 //
 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
+constexpr const StringLiteral
+    mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
 
 namespace mlir {
 namespace impl {
@@ -273,6 +275,12 @@
   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 @@
                             << DLTIDialect::kDataLayoutEndiannessBig << "' or '"
                             << DLTIDialect::kDataLayoutEndiannessLittle << "'";
     }
+    if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey)
+      return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;
   }
 };
Index: mlir/lib/Interfaces/DataLayoutInterfaces.cpp
===================================================================
--- mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -213,6 +213,18 @@
   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 @@
 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(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 @@
   });
 }
 
+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
 //===----------------------------------------------------------------------===//
Index: mlir/lib/Target/LLVMIR/ModuleImport.cpp
===================================================================
--- mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ 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 @@
   // 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 @@
           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);
+      }
     }
   }
 
Index: mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
===================================================================
--- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -79,6 +79,15 @@
       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();
   }
Index: mlir/test/Dialect/LLVMIR/layout.mlir
===================================================================
--- mlir/test/Dialect/LLVMIR/layout.mlir
+++ mlir/test/Dialect/LLVMIR/layout.mlir
@@ -4,37 +4,44 @@
   // 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 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>
Index: mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
===================================================================
--- mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -40,11 +40,18 @@
       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)
+
+          });
     });
   }
 };
Index: mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
===================================================================
--- mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -22,6 +22,8 @@
 
 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 @@
   }
   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 @@
     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 @@
   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 @@
   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 @@
   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 @@
   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 @@
   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) {