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

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 @@ public:
   /// 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 @@ private:
 
   /// 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 @@ public:
         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 @@ private:
   /// 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"

mlir/test/Target/LLVMIR/Import/data-layout.ll

@@ -30,7 +30,8 @@ target datalayout = ""
 ; CHECK-SAME:   "dlti.endianness" = "little"
 ; CHECK-SAME:   "dlti.mangling_mode" = "e"
 ; CHECK-SAME:   "dlti.stack_alignment" = 128 : i64
-target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+; CHECK-SAME:   "dlti.function_pointer_alignment" = #dlti.function_pointer_alignment<32, function_dependent = true>
+target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-Fn32"
 
 ; // -----
 

mlir/test/Target/LLVMIR/data-layout.mlir

@@ -5,6 +5,7 @@
 // CHECK: A4-
 // CHECK: S128-
 // CHECK: m:e-
+// CHECK: Fn32
 // CHECK: i64:64:128
 // CHECK: f80:128:256
 // CHECK: p0:32:64:128:32
@@ -14,6 +15,8 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
 #dlti.dl_entry<"dlti.alloca_memory_space", 4 : ui32>,
 #dlti.dl_entry<"dlti.stack_alignment", 128 : i32>,
 #dlti.dl_entry<"dlti.mangling_mode", "e">,
+#dlti.dl_entry<"dlti.function_pointer_alignment",
+               #dlti.function_pointer_alignment<32, function_dependent = true>>,
 #dlti.dl_entry<index, 64>,
 #dlti.dl_entry<i64, dense<[64,128]> : vector<2xi64>>,
 #dlti.dl_entry<f80, dense<[128,256]> : vector<2xi64>>,
@@ -29,12 +32,15 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
 
 // CHECK: target datalayout
 // CHECK: e
+// CHECK: Fi64
 // CHECK-NOT: A0
 // CHECK-NOT: S0
 module attributes {dlti.dl_spec = #dlti.dl_spec<
 #dlti.dl_entry<"dlti.endianness", "little">,
 #dlti.dl_entry<"dlti.alloca_memory_space", 0 : ui32>,
-#dlti.dl_entry<"dlti.stack_alignment", 0 : i32>
+#dlti.dl_entry<"dlti.stack_alignment", 0 : i32>,
+#dlti.dl_entry<"dlti.function_pointer_alignment",
+               #dlti.function_pointer_alignment<64, function_dependent = false>>
 >} {
   llvm.func @bar() {
     llvm.return

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

@@ -48,6 +48,7 @@ struct TestDataLayoutQuery
       Attribute programMemorySpace = layout.getProgramMemorySpace();
       Attribute globalMemorySpace = layout.getGlobalMemorySpace();
       uint64_t stackAlignment = layout.getStackAlignment();
+      Attribute functionPointerAlignment = layout.getFunctionPointerAlignment();
 
       auto convertTypeSizeToAttr = [&](llvm::TypeSize typeSize) -> Attribute {
         if (!typeSize.isScalable())
@@ -90,7 +91,16 @@ struct TestDataLayoutQuery
                                     ? builder.getUI32IntegerAttr(0)
                                     : globalMemorySpace),
            builder.getNamedAttr("stack_alignment",
-                                builder.getIndexAttr(stackAlignment))});
+                                builder.getIndexAttr(stackAlignment)),
+           builder.getNamedAttr("function_pointer_alignment",
+                                functionPointerAlignment == Attribute()
+                                    ? FunctionPointerAlignmentAttr::get(
+                                          builder.getContext(), 0,
+                                          /*function_dependent=*/false)
+                                    : functionPointerAlignment)
+
+          });
+
     });
   }
 };

mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

@@ -35,6 +35,8 @@ constexpr static llvm::StringLiteral kGlobalKeyName =
     "dltest.global_memory_space";
 constexpr static llvm::StringLiteral kStackAlignmentKeyName =
     "dltest.stack_alignment";
+constexpr static llvm::StringLiteral kFunctionPointerAlignmentKeyName =
+    "dltest.function_pointer_alignment";
 
 constexpr static llvm::StringLiteral kTargetSystemDescAttrName =
     "dl_target_sys_desc_test.target_system_spec";
@@ -102,6 +104,9 @@ struct CustomDataLayoutSpec
   StringAttr getStackAlignmentIdentifier(MLIRContext *context) const {
     return Builder(context).getStringAttr(kStackAlignmentKeyName);
   }
+  StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const {
+    return Builder(context).getStringAttr(kFunctionPointerAlignmentKeyName);
+  }
   FailureOr<Attribute> query(DataLayoutEntryKey key) const {
     return llvm::cast<mlir::DataLayoutSpecInterface>(*this).queryHelper(key);
   }
@@ -494,6 +499,7 @@ module {}
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
+  EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
   EXPECT_EQ(layout.getManglingMode(), Attribute());
 }
 
@@ -571,6 +577,7 @@ TEST(DataLayout, EmptySpec) {
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
   EXPECT_EQ(layout.getManglingMode(), Attribute());
+  EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
 
   EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
@@ -594,7 +601,9 @@ TEST(DataLayout, SpecWithEntries) {
   #dlti.dl_entry<"dltest.program_memory_space", 3 : i32>,
   #dlti.dl_entry<"dltest.global_memory_space", 2 : i32>,
   #dlti.dl_entry<"dltest.stack_alignment", 128 : i32>,
-  #dlti.dl_entry<"dltest.mangling_mode", "o">
+  #dlti.dl_entry<"dltest.mangling_mode", "o">,
+  #dlti.dl_entry<"dltest.function_pointer_alignment",
+                 #dlti.function_pointer_alignment<64, function_dependent = true>>
 > } : () -> ()
   )MLIR";
 
@@ -633,6 +642,9 @@ TEST(DataLayout, SpecWithEntries) {
   EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2));
   EXPECT_EQ(layout.getStackAlignment(), 128u);
   EXPECT_EQ(layout.getManglingMode(), Builder(&ctx).getStringAttr("o"));
+  EXPECT_EQ(
+      layout.getFunctionPointerAlignment(),
+      FunctionPointerAlignmentAttr::get(&ctx, 64, /*function_dependent=*/true));
 }
 
 TEST(DataLayout, SpecWithTargetSystemDescEntries) {