[MLIR][LLVM][DLTI] Handle data layout token 'Fn32' (#141167)

Author: bcardosolopes

mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td

@@ -92,6 +92,9 @@ def DLTI_DataLayoutSpecAttr :
     /// Returns the stack alignment identifier.
     StringAttr getStackAlignmentIdentifier(MLIRContext *context) const;
 
+    /// Returns the function pointer alignment identifier.
+    StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const;
+
     /// Returns the attribute associated with the key.
     FailureOr<Attribute> query(DataLayoutEntryKey key) {
       return ::llvm::cast<mlir::DataLayoutSpecInterface>(*this).queryHelper(key);
@@ -249,4 +252,35 @@ def DLTI_TargetDeviceSpecAttr :
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// FunctionPointerAlignmentAttr
+//===----------------------------------------------------------------------===//
+
+def DLTI_FunctionPointerAlignmentAttr :
+    DLTIAttr<"FunctionPointerAlignment", []> {
+  let summary = "An attribute to represent function pointer alignment.";
+  let description = [{
+    Function pointer alignment specifies the minimum alignment of function
+    pointers, it's a multiple of `alignment`. This alignment can also depend
+    on the target function, indicated by `function_dependent`.
+    Example:
+    ```
+    #dlti.dl_entry<"dlti.function_pointer_alignment",
+               #dlti.function_pointer_alignment<64, function_dependent = false>>
+    ```
+  }];
+  let parameters = (ins
+    "uint64_t":$alignment,
+    "bool":$function_dependent
+  );
+  let assemblyFormat = [{
+    `<`
+    $alignment `,`
+    `function_dependent` `=` $function_dependent
+    `>`
+  }];
+  let mnemonic = "function_pointer_alignment";
+}
+
+
 #endif  // MLIR_DIALECT_DLTI_DLTIATTRS_TD

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

@@ -58,7 +58,7 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
-    
+
     constexpr const static ::llvm::StringLiteral
     kDataLayoutProgramMemorySpaceKey = "dlti.program_memory_space";
 
@@ -67,6 +67,9 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
+
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutFunctionPointerAlignmentKey = "dlti.function_pointer_alignment";
   }];
 
   let useDefaultAttributePrinterParser = 1;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -101,6 +101,10 @@ Attribute getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry);
 /// DataLayoutInterface if specified, otherwise returns the default.
 uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 
+/// Default handler for the function pointer alignment request. Dispatches to
+/// the DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultFunctionPointerAlignment(DataLayoutEntryInterface entry);
+
 /// Returns the value of the property from the specified DataLayoutEntry. If the
 /// property is missing from the entry, returns std::nullopt.
 std::optional<Attribute> getDevicePropertyValue(DataLayoutEntryInterface entry);
@@ -259,6 +263,9 @@ class DataLayout {
   /// unspecified.
   uint64_t getStackAlignment() const;
 
+  /// Returns function pointer alignment.
+  Attribute getFunctionPointerAlignment() const;
+
   /// Returns the value of the specified property if the property is defined for
   /// the given device ID, otherwise returns std::nullopt.
   std::optional<Attribute>
@@ -303,6 +310,8 @@ class DataLayout {
 
   /// Cache for stack alignment.
   mutable std::optional<uint64_t> stackAlignment;
+  /// Cache for function pointer alignment.
+  mutable std::optional<Attribute> functionPointerAlignment;
 };
 
 } // namespace mlir

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -171,6 +171,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
       /*methodName=*/"getStackAlignmentIdentifier",
       /*args=*/(ins "::mlir::MLIRContext *":$context)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the function pointer alignment identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getFunctionPointerAlignmentIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     // Implementations may override this if they have an efficient lookup
     // mechanism.
     InterfaceMethod<
@@ -553,6 +559,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         return ::mlir::detail::getDefaultStackAlignment(entry);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the function pointer alignment in bits computed "
+                      "using the relevant entries. The data layout object "
+                      "can be used for recursive queries.",
+      /*retTy=*/"Attribute",
+      /*methodName=*/"getFunctionPointerAlignment",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultFunctionPointerAlignment(entry);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the value of the property, if the property is "
                       "defined. Otherwise, it returns std::nullopt.",

mlir/lib/Dialect/DLTI/DLTI.cpp

@@ -420,6 +420,12 @@ DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const {
       DLTIDialect::kDataLayoutStackAlignmentKey);
 }
 
+StringAttr DataLayoutSpecAttr::getFunctionPointerAlignmentIdentifier(
+    MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutFunctionPointerAlignmentKey);
+}
+
 /// Parses an attribute with syntax:
 ///   dl-spec-attr ::= `#dlti.` `dl_spec` `<` entry-list `>`
 ///   entry-list ::= | entry | entry `,` entry-list
@@ -625,6 +631,7 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
         entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||
+        entryName == DLTIDialect::kDataLayoutFunctionPointerAlignmentKey ||
         entryName == DLTIDialect::kDataLayoutManglingModeKey)
       return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

@@ -312,6 +312,15 @@ mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
   return value.getValue().getZExtValue();
 }
 
+// Returns the function pointer alignment if specified in the given entry. If
+// the entry is empty the default alignment zero is returned.
+Attribute mlir::detail::getDefaultFunctionPointerAlignment(
+    DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface())
+    return Attribute();
+  return entry.getValue();
+}
+
 std::optional<Attribute>
 mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
   if (entry == DataLayoutEntryInterface())
@@ -710,6 +719,23 @@ uint64_t mlir::DataLayout::getStackAlignment() const {
   return *stackAlignment;
 }
 
+Attribute mlir::DataLayout::getFunctionPointerAlignment() const {
+  checkValid();
+  if (functionPointerAlignment)
+    return *functionPointerAlignment;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getFunctionPointerAlignmentIdentifier(
+            originalLayout.getContext()));
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    functionPointerAlignment = iface.getFunctionPointerAlignment(entry);
+  else
+    functionPointerAlignment =
+        detail::getDefaultFunctionPointerAlignment(entry);
+  return *functionPointerAlignment;
+}
+
 std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
     TargetSystemSpecInterface::DeviceID deviceID,
     StringAttr propertyName) const {

mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp

@@ -221,6 +221,36 @@ DataLayoutImporter::tryToEmplaceStackAlignmentEntry(StringRef token) {
   return success();
 }
 
+LogicalResult DataLayoutImporter::tryToEmplaceFunctionPointerAlignmentEntry(
+    StringRef fnPtrString, StringRef token) {
+  auto key = StringAttr::get(
+      context, DLTIDialect::kDataLayoutFunctionPointerAlignmentKey);
+  if (keyEntries.count(key))
+    return success();
+
+  // The data layout entry for "F<type><abi>". <abi> is the aligment value,
+  // preceded by one of the two possible <types>:
+  // "i": The alignment of function pointers is independent of the alignment of
+  //      functions, and is a multiple of <abi>.
+  // "n": The alignment of function pointers is a multiple of the explicit
+  //      alignment specified on the function, and is a multiple of <abi>.
+  bool functionDependent = false;
+  if (fnPtrString == "n")
+    functionDependent = true;
+  else if (fnPtrString != "i")
+    return failure();
+
+  FailureOr<uint64_t> alignment = tryToParseInt(token);
+  if (failed(alignment))
+    return failure();
+
+  keyEntries.try_emplace(
+      key, DataLayoutEntryAttr::get(
+               key, FunctionPointerAlignmentAttr::get(
+                        key.getContext(), *alignment, functionDependent)));
+  return success();
+}
+
 void DataLayoutImporter::translateDataLayout(
     const llvm::DataLayout &llvmDataLayout) {
   dataLayout = {};
@@ -330,6 +360,14 @@ void DataLayoutImporter::translateDataLayout(
         return;
       continue;
     }
+    // Parse function pointer alignment specifications.
+    // Note that prefix here is "Fn" or "Fi", not a single character.
+    if (prefix->starts_with("F")) {
+      StringRef nextPrefix = prefix->drop_front(1);
+      if (failed(tryToEmplaceFunctionPointerAlignmentEntry(nextPrefix, token)))
+        return;
+      continue;
+    }
 
     // Store all tokens that have not been handled.
     unhandledTokens.push_back(lastToken);

mlir/lib/Target/LLVMIR/DataLayoutImporter.h

@@ -108,6 +108,11 @@ class DataLayoutImporter {
   /// Adds a stack alignment entry if there is none yet.
   LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token);
 
+  /// Adds a function pointer alignment entry if there is none yet.
+  LogicalResult
+  tryToEmplaceFunctionPointerAlignmentEntry(StringRef fnPtrAlignEntry,
+                                            StringRef token);
+
   std::string layoutStr = {};
   StringRef lastToken = {};
   SmallVector<StringRef> unhandledTokens;

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -240,6 +240,16 @@ translateDataLayout(DataLayoutSpecInterface attribute,
       layoutStream << "-S" << alignment;
       continue;
     }
+    if (key.getValue() == DLTIDialect::kDataLayoutFunctionPointerAlignmentKey) {
+      auto value = cast<FunctionPointerAlignmentAttr>(entry.getValue());
+      uint64_t alignment = value.getAlignment();
+      // Skip the default function pointer alignment.
+      if (alignment == 0)
+        continue;
+      layoutStream << "-F" << (value.getFunctionDependent() ? "n" : "i")
+                   << alignment;
+      continue;
+    }
     emitError(*loc) << "unsupported data layout key " << key;
     return failure();
   }

mlir/test/Dialect/LLVMIR/layout.mlir

@@ -8,6 +8,8 @@ module {
     // CHECK: bitsize = 64
     // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
     // CHECK: mangling_mode = ""
@@ -21,6 +23,8 @@ module {
     // CHECK: bitsize = 64
     // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
     // CHECK: mangling_mode = ""
@@ -34,6 +38,8 @@ module {
     // CHECK: bitsize = 64
     // CHECK: default_memory_space = 0
     // CHECK: endianness = ""
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: #dlti.function_pointer_alignment<0, function_dependent = false>,
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
     // CHECK: mangling_mode = ""
@@ -58,7 +64,9 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
   #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
   #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>,
-  #dlti.dl_entry<"dlti.mangling_mode", "e">
+  #dlti.dl_entry<"dlti.mangling_mode", "e">,
+  #dlti.dl_entry<"dlti.function_pointer_alignment",
+                 "#dlti.function_pointer_alignment<32, function_dependent = true>">
 >} {
   // CHECK: @spec
   func.func @spec() {
@@ -67,6 +75,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: bitsize = 32
     // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
     // CHECK: mangling_mode = "e"
@@ -80,6 +90,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: bitsize = 32
     // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
     // CHECK: preferred = 8
@@ -92,6 +104,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: bitsize = 64
     // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 64
     // CHECK: mangling_mode = "e"
@@ -105,6 +119,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: bitsize = 32
     // CHECK: default_memory_space = 7
     // CHECK: endianness = "little"
+    // CHECK: function_pointer_alignment =
+    // CHECK-SAME: "#dlti.function_pointer_alignment<32, function_dependent = true>",
     // CHECK: global_memory_space = 2
     // CHECK: index = 24
     // CHECK: mangling_mode = "e"