[mlir] Add global and program memory space handling to the data layout subsystem

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

@@ -103,6 +103,12 @@ class DataLayoutSpecAttr
   /// Returns the alloca memory space identifier.
   StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
 
+  /// Returns the program memory space identifier.
+  StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const;
+
+  /// Returns the global memory space identifier.
+  StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const;
+
   /// Returns the stack alignment identifier.
   StringAttr getStackAlignmentIdentifier(MLIRContext *context) const;
 

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

@@ -39,6 +39,12 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutProgramMemorySpaceKey = "dlti.program_memory_space";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutGlobalMemorySpaceKey = "dlti.global_memory_space";
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutStackAlignmentKey = "dlti.stack_alignment";

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -61,6 +61,14 @@ getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
 
+/// Default handler for program memory space request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultProgramMemorySpace(DataLayoutEntryInterface entry);
+
+/// Default handler for global memory space request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry);
+
 /// Default handler for the stack alignment request. Dispatches to the
 /// DataLayoutInterface if specified, otherwise returns the default.
 uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
@@ -175,6 +183,12 @@ class DataLayout {
   /// Returns the memory space used for AllocaOps.
   Attribute getAllocaMemorySpace() const;
 
+  /// Returns the memory space used for program memory operations.
+  Attribute getProgramMemorySpace() const;
+
+  /// Returns the memory space used for global operations.
+  Attribute getGlobalMemorySpace() const;
+
   /// Returns the natural alignment of the stack in bits. Alignment promotion of
   /// stack variables should be limited to the natural stack alignment to
   /// prevent dynamic stack alignment. Returns zero if the stack alignment is
@@ -203,8 +217,10 @@ class DataLayout {
   mutable DenseMap<Type, uint64_t> abiAlignments;
   mutable DenseMap<Type, uint64_t> preferredAlignments;
 
-  /// Cache for alloca memory space.
+  /// Cache for alloca, global, and program memory spaces.
   mutable std::optional<Attribute> allocaMemorySpace;
+  mutable std::optional<Attribute> programMemorySpace;
+  mutable std::optional<Attribute> globalMemorySpace;
 
   /// Cache for stack alignment.
   mutable std::optional<uint64_t> stackAlignment;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -112,6 +112,18 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> {
       /*methodName=*/"getAllocaMemorySpaceIdentifier",
       /*args=*/(ins "::mlir::MLIRContext *":$context)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the program memory space identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getProgramMemorySpaceIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
+    InterfaceMethod<
+      /*description=*/"Returns the global memory space identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getGlobalMemorySpaceIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     InterfaceMethod<
       /*description=*/"Returns the stack alignment identifier.",
       /*retTy=*/"::mlir::StringAttr",
@@ -280,6 +292,30 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         return ::mlir::detail::getDefaultAllocaMemorySpace(entry);
       }]
     >,
+    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=*/"getProgramMemorySpace",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultProgramMemorySpace(entry);
+      }]
+    >,
+    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=*/"getGlobalMemorySpace",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultGlobalMemorySpace(entry);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the natural stack alignment in bits computed "
                       "using the relevant entries. The data layout object "

mlir/lib/Dialect/DLTI/DLTI.cpp

@@ -108,6 +108,11 @@ void DataLayoutEntryAttr::print(AsmPrinter &os) const {
 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
 constexpr const StringLiteral
     mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
+constexpr const StringLiteral
+    mlir::DLTIDialect::kDataLayoutProgramMemorySpaceKey;
+constexpr const StringLiteral
+    mlir::DLTIDialect::kDataLayoutGlobalMemorySpaceKey;
+
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutStackAlignmentKey;
 
 namespace mlir {
@@ -282,6 +287,17 @@ DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
       DLTIDialect::kDataLayoutAllocaMemorySpaceKey);
 }
 
+StringAttr DataLayoutSpecAttr::getProgramMemorySpaceIdentifier(
+    MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutProgramMemorySpaceKey);
+}
+
+StringAttr
+DataLayoutSpecAttr::getGlobalMemorySpaceIdentifier(MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutGlobalMemorySpaceKey);
+}
 StringAttr
 DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const {
   return Builder(context).getStringAttr(
@@ -345,6 +361,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
                             << DLTIDialect::kDataLayoutEndiannessLittle << "'";
     }
     if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
+        entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey ||
+        entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutStackAlignmentKey)
       return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -230,6 +230,30 @@ mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
   return entry.getValue();
 }
 
+// Returns the memory space used for the program memory space.  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::getDefaultProgramMemorySpace(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface()) {
+    return Attribute();
+  }
+
+  return entry.getValue();
+}
+
+// Returns the memory space used for global the global memory space. 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::getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface()) {
+    return Attribute();
+  }
+
+  return entry.getValue();
+}
+
 // Returns the stack alignment if specified in the given entry. If the entry is
 // empty the default alignment zero is returned.
 uint64_t
@@ -382,7 +406,8 @@ mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
 
 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
     : originalLayout(getCombinedDataLayout(op)), scope(op),
-      allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
+      allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt),
+      globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
@@ -391,7 +416,8 @@ mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
 
 mlir::DataLayout::DataLayout(ModuleOp op)
     : originalLayout(getCombinedDataLayout(op)), scope(op),
-      allocaMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
+      allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt),
+      globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   checkMissingLayout(originalLayout, op);
   collectParentLayouts(op, layoutStack);
@@ -510,6 +536,38 @@ mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
   return *allocaMemorySpace;
 }
 
+mlir::Attribute mlir::DataLayout::getProgramMemorySpace() const {
+  checkValid();
+  if (programMemorySpace)
+    return *programMemorySpace;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getProgramMemorySpaceIdentifier(
+            originalLayout.getContext()));
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    programMemorySpace = iface.getProgramMemorySpace(entry);
+  else
+    programMemorySpace = detail::getDefaultProgramMemorySpace(entry);
+  return *programMemorySpace;
+}
+
+mlir::Attribute mlir::DataLayout::getGlobalMemorySpace() const {
+  checkValid();
+  if (globalMemorySpace)
+    return *globalMemorySpace;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getGlobalMemorySpaceIdentifier(
+            originalLayout.getContext()));
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    globalMemorySpace = iface.getGlobalMemorySpace(entry);
+  else
+    globalMemorySpace = detail::getDefaultGlobalMemorySpace(entry);
+  return *globalMemorySpace;
+}
+
 uint64_t mlir::DataLayout::getStackAlignment() const {
   checkValid();
   if (stackAlignment)

mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp

@@ -164,9 +164,9 @@ DataLayoutImporter::tryToEmplaceEndiannessEntry(StringRef endianness,
 }
 
 LogicalResult
-DataLayoutImporter::tryToEmplaceAllocaAddrSpaceEntry(StringRef token) {
-  auto key =
-      StringAttr::get(context, DLTIDialect::kDataLayoutAllocaMemorySpaceKey);
+DataLayoutImporter::tryToEmplaceAddrSpaceEntry(StringRef token,
+                                               llvm::StringLiteral spaceKey) {
+  auto key = StringAttr::get(context, spaceKey);
   if (keyEntries.count(key))
     return success();
 
@@ -247,9 +247,24 @@ void DataLayoutImporter::translateDataLayout(
         return;
       continue;
     }
+    // Parse the program address space.
+    if (*prefix == "P") {
+      if (failed(tryToEmplaceAddrSpaceEntry(
+              token, DLTIDialect::kDataLayoutProgramMemorySpaceKey)))
+        return;
+      continue;
+    }
+    // Parse the global address space.
+    if (*prefix == "G") {
+      if (failed(tryToEmplaceAddrSpaceEntry(
+              token, DLTIDialect::kDataLayoutGlobalMemorySpaceKey)))
+        return;
+      continue;
+    }
     // Parse the alloca address space.
     if (*prefix == "A") {
-      if (failed(tryToEmplaceAllocaAddrSpaceEntry(token)))
+      if (failed(tryToEmplaceAddrSpaceEntry(
+              token, DLTIDialect::kDataLayoutAllocaMemorySpaceKey)))
         return;
       continue;
     }

mlir/lib/Target/LLVMIR/DataLayoutImporter.h

@@ -97,7 +97,8 @@ class DataLayoutImporter {
                                             StringRef token);
 
   /// Adds an alloca address space entry if there is none yet.
-  LogicalResult tryToEmplaceAllocaAddrSpaceEntry(StringRef token);
+  LogicalResult tryToEmplaceAddrSpaceEntry(StringRef token,
+                                           llvm::StringLiteral spaceKey);
 
   /// Adds a stack alignment entry if there is none yet.
   LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -190,6 +190,26 @@ translateDataLayout(DataLayoutSpecInterface attribute,
       layoutStream.flush();
       continue;
     }
+    if (key.getValue() == DLTIDialect::kDataLayoutProgramMemorySpaceKey) {
+      auto value = cast<IntegerAttr>(entry.getValue());
+      uint64_t space = value.getValue().getZExtValue();
+      // Skip the default address space.
+      if (space == 0)
+        continue;
+      layoutStream << "-P" << space;
+      layoutStream.flush();
+      continue;
+    }
+    if (key.getValue() == DLTIDialect::kDataLayoutGlobalMemorySpaceKey) {
+      auto value = cast<IntegerAttr>(entry.getValue());
+      uint64_t space = value.getValue().getZExtValue();
+      // Skip the default address space.
+      if (space == 0)
+        continue;
+      layoutStream << "-G" << space;
+      layoutStream.flush();
+      continue;
+    }
     if (key.getValue() == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) {
       auto value = cast<IntegerAttr>(entry.getValue());
       uint64_t space = value.getValue().getZExtValue();