[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

Author: jsjodin

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);
 

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;

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

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); }];

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;
   }
 };

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,15 +358,17 @@ 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);
 #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);
@@ -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
 //===----------------------------------------------------------------------===//

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);
+      }
     }
   }
 

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();
   }

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>

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)
+
+          });
     });
   }
 };