[MLIR][DLTI] Add mangling style (#125875)

ghehg · web-flow · commit 8c0e9adc5c1f · 2025-03-05T13:47:50.000+01:00
Add mangling style as a spec entry to datalayout, and implemented importing and exporting LLVM IR to MLIR (LLVM dialect). Its represented as string as the scope of this PR is to preserve info from LLVM IR, so client in MLIR still need to map deduce the meaning of the string, like "e" means ELF, "o" for Mach-O, etc. it addresses one of issues mentioned in this [issue](#126046)
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
@@ -77,6 +77,9 @@ def DLTI_DataLayoutSpecAttr :
     /// Returns the alloca memory space identifier.
     StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
 
+    /// Returns the mangling mode identifier.
+    StringAttr getManglingModeIdentifier(MLIRContext *context) const;
+
     /// Returns the program memory space identifier.
     StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const;
 
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td
@@ -49,6 +49,10 @@ def DLTI_Dialect : Dialect {
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessLittle = "little";
 
+    // Mangling mode, a.k.a mangling style used to mangle llvm names.
+    constexpr const static ::llvm::StringLiteral
+    kDataLayoutManglingModeKey = "dlti.mangling_mode";
+
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
     
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -78,6 +78,10 @@ Attribute getDefaultEndianness(DataLayoutEntryInterface entry);
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
 
+/// Default handler for mangling mode request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultManglingMode(DataLayoutEntryInterface entry);
+
 /// Default handler for program memory space request. Dispatches to the
 /// DataLayoutInterface if specified, otherwise returns the default.
 Attribute getDefaultProgramMemorySpace(DataLayoutEntryInterface entry);
@@ -230,6 +234,9 @@ class DataLayout {
   /// Returns the memory space used for AllocaOps.
   Attribute getAllocaMemorySpace() const;
 
+  /// Returns the mangling mode.
+  Attribute getManglingMode() const;
+
   /// Returns the memory space used for program memory operations.
   Attribute getProgramMemorySpace() const;
 
@@ -276,6 +283,8 @@ class DataLayout {
 
   /// Cache for the endianness.
   mutable std::optional<Attribute> endianness;
+  /// Cache for the mangling mode.
+  mutable std::optional<Attribute> manglingMode;
   /// Cache for alloca, global, and program memory spaces.
   mutable std::optional<Attribute> allocaMemorySpace;
   mutable std::optional<Attribute> programMemorySpace;
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -153,6 +153,12 @@ def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface", [DLTIQuer
       /*methodName=*/"getGlobalMemorySpaceIdentifier",
       /*args=*/(ins "::mlir::MLIRContext *":$context)
     >,
+    InterfaceMethod<
+      /*description=*/"Returns the mangling mode identifier.",
+      /*retTy=*/"::mlir::StringAttr",
+      /*methodName=*/"getManglingModeIdentifier",
+      /*args=*/(ins "::mlir::MLIRContext *":$context)
+    >,
     InterfaceMethod<
       /*description=*/"Returns the stack alignment identifier.",
       /*retTy=*/"::mlir::StringAttr",
@@ -481,6 +487,18 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
         return ::mlir::detail::getDefaultAllocaMemorySpace(entry);
       }]
     >,
+    StaticInterfaceMethod<
+      /*description=*/"Returns the mangling mode computed "
+                      "using the relevant entries. The data layout object "
+                      "can be used for recursive queries.",
+      /*retTy=*/"::mlir::Attribute",
+      /*methodName=*/"getManglingMode",
+      /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::detail::getDefaultManglingMode(entry);
+      }]
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the memory space used by the ABI computed "
                       "using the relevant entries. The data layout object "
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -397,6 +397,12 @@ DataLayoutSpecAttr::getGlobalMemorySpaceIdentifier(MLIRContext *context) const {
       DLTIDialect::kDataLayoutGlobalMemorySpaceKey);
 }
 
+StringAttr
+DataLayoutSpecAttr::getManglingModeIdentifier(MLIRContext *context) const {
+  return Builder(context).getStringAttr(
+      DLTIDialect::kDataLayoutManglingModeKey);
+}
+
 StringAttr
 DataLayoutSpecAttr::getStackAlignmentIdentifier(MLIRContext *context) const {
   return Builder(context).getStringAttr(
@@ -606,7 +612,8 @@ class TargetDataLayoutInterface : public DataLayoutDialectInterface {
     if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutProgramMemorySpaceKey ||
         entryName == DLTIDialect::kDataLayoutGlobalMemorySpaceKey ||
-        entryName == DLTIDialect::kDataLayoutStackAlignmentKey)
+        entryName == DLTIDialect::kDataLayoutStackAlignmentKey ||
+        entryName == DLTIDialect::kDataLayoutManglingModeKey)
       return success();
     return emitError(loc) << "unknown data layout entry name: " << entryName;
   }
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -258,6 +258,15 @@ mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
   return entry.getValue();
 }
 
+// Returns the mangling mode if specified in the given entry.
+// If the entry is empty, an empty attribute is returned.
+Attribute mlir::detail::getDefaultManglingMode(DataLayoutEntryInterface entry) {
+  if (entry == DataLayoutEntryInterface())
+    return Attribute();
+
+  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.
@@ -612,6 +621,22 @@ mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
   return *allocaMemorySpace;
 }
 
+mlir::Attribute mlir::DataLayout::getManglingMode() const {
+  checkValid();
+  if (manglingMode)
+    return *manglingMode;
+  DataLayoutEntryInterface entry;
+  if (originalLayout)
+    entry = originalLayout.getSpecForIdentifier(
+        originalLayout.getManglingModeIdentifier(originalLayout.getContext()));
+
+  if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+    manglingMode = iface.getManglingMode(entry);
+  else
+    manglingMode = detail::getDefaultManglingMode(entry);
+  return *manglingMode;
+}
+
 mlir::Attribute mlir::DataLayout::getProgramMemorySpace() const {
   checkValid();
   if (programMemorySpace)
diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp b/mlir/lib/Target/LLVMIR/DataLayoutImporter.cpp
@@ -163,6 +163,21 @@ DataLayoutImporter::tryToEmplaceEndiannessEntry(StringRef endianness,
   return success();
 }
 
+LogicalResult DataLayoutImporter::tryToEmplaceManglingModeEntry(
+    StringRef token, llvm::StringLiteral manglingKey) {
+  auto key = StringAttr::get(context, manglingKey);
+  if (keyEntries.count(key))
+    return success();
+
+  token.consume_front(":");
+  if (token.empty())
+    return failure();
+
+  keyEntries.try_emplace(
+      key, DataLayoutEntryAttr::get(key, StringAttr::get(context, token)));
+  return success();
+}
+
 LogicalResult
 DataLayoutImporter::tryToEmplaceAddrSpaceEntry(StringRef token,
                                                llvm::StringLiteral spaceKey) {
@@ -254,6 +269,13 @@ void DataLayoutImporter::translateDataLayout(
         return;
       continue;
     }
+    // Parse the mangling mode.
+    if (*prefix == "m") {
+      if (failed(tryToEmplaceManglingModeEntry(
+              token, DLTIDialect::kDataLayoutManglingModeKey)))
+        return;
+      continue;
+    }
     // Parse the global address space.
     if (*prefix == "G") {
       if (failed(tryToEmplaceAddrSpaceEntry(
diff --git a/mlir/lib/Target/LLVMIR/DataLayoutImporter.h b/mlir/lib/Target/LLVMIR/DataLayoutImporter.h
@@ -100,6 +100,10 @@ class DataLayoutImporter {
   LogicalResult tryToEmplaceAddrSpaceEntry(StringRef token,
                                            llvm::StringLiteral spaceKey);
 
+  /// Adds an mangling mode entry if there is none yet.
+  LogicalResult tryToEmplaceManglingModeEntry(StringRef token,
+                                              llvm::StringLiteral manglingKey);
+
   /// Adds a stack alignment entry if there is none yet.
   LogicalResult tryToEmplaceStackAlignmentEntry(StringRef token);
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -199,6 +199,11 @@ translateDataLayout(DataLayoutSpecInterface attribute,
       layoutStream << "-" << (isLittleEndian ? "e" : "E");
       continue;
     }
+    if (key.getValue() == DLTIDialect::kDataLayoutManglingModeKey) {
+      auto value = cast<StringAttr>(entry.getValue());
+      layoutStream << "-m:" << value.getValue();
+      continue;
+    }
     if (key.getValue() == DLTIDialect::kDataLayoutProgramMemorySpaceKey) {
       auto value = cast<IntegerAttr>(entry.getValue());
       uint64_t space = value.getValue().getZExtValue();
diff --git a/mlir/test/Dialect/LLVMIR/layout.mlir b/mlir/test/Dialect/LLVMIR/layout.mlir
@@ -9,6 +9,7 @@ module {
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -20,6 +21,7 @@ module {
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -31,6 +33,7 @@ module {
     // CHECK: endianness = ""
     // CHECK: global_memory_space = 0
     // CHECK: index = 64
+    // CHECK: mangling_mode = ""
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 0
     // CHECK: size = 8
@@ -50,7 +53,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
   #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
   #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.stack_alignment", 128 : i64>,
+  #dlti.dl_entry<"dlti.mangling_mode", "e">
 >} {
   // CHECK: @spec
   func.func @spec() {
@@ -60,6 +64,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 32
+    // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
@@ -82,6 +87,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 64
+    // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 8
@@ -93,6 +99,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: endianness = "little"
     // CHECK: global_memory_space = 2
     // CHECK: index = 24
+    // CHECK: mangling_mode = "e"
     // CHECK: preferred = 8
     // CHECK: program_memory_space = 3
     // CHECK: size = 4
diff --git a/mlir/test/Target/LLVMIR/Import/data-layout.ll b/mlir/test/Target/LLVMIR/Import/data-layout.ll
@@ -29,6 +29,7 @@ target datalayout = ""
 ; CHECK-DAG:   !llvm.ptr<271> = dense<32> : vector<4xi64>
 ; CHECK-DAG:   !llvm.ptr<272> = dense<64> : vector<4xi64>
 ; CHECK-DAG:   "dlti.stack_alignment" = 128 : i64
+; CHECK-DAG:   "dlti.mangling_mode" = "e"
 target datalayout = "e-m:e-p270:32:64-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
 ; // -----
diff --git a/mlir/test/Target/LLVMIR/Import/import-failure.ll b/mlir/test/Target/LLVMIR/Import/import-failure.ll
@@ -344,6 +344,12 @@ target datalayout = "e-ni:42-i64:64"
 
 ; // -----
 
+; CHECK:      import-failure.ll
+; CHECK-SAME: malformed specification, must be of the form "m:<mangling>"
+target datalayout = "e-m-i64:64"
+
+; // -----
+
 ; CHECK:      <unknown>
 ; CHECK-SAME: incompatible call and callee types: '!llvm.func<void (i64)>' and '!llvm.func<void (ptr)>'
 define void @incompatible_call_and_callee_types() {
diff --git a/mlir/test/Target/LLVMIR/data-layout.mlir b/mlir/test/Target/LLVMIR/data-layout.mlir
@@ -4,6 +4,7 @@
 // CHECK: E-
 // CHECK: A4-
 // CHECK: S128-
+// CHECK: m:e-
 // CHECK: i64:64:128
 // CHECK: f80:128:256
 // CHECK: p0:32:64:128:32
@@ -12,6 +13,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<
 #dlti.dl_entry<"dlti.endianness", "big">,
 #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<index, 64>,
 #dlti.dl_entry<i64, dense<[64,128]> : vector<2xi64>>,
 #dlti.dl_entry<f80, dense<[128,256]> : vector<2xi64>>,
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -43,6 +43,7 @@ struct TestDataLayoutQuery
       uint64_t index = layout.getTypeIndexBitwidth(op.getType()).value_or(0);
       Attribute endianness = layout.getEndianness();
       Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
+      Attribute manglingMode = layout.getManglingMode();
       Attribute programMemorySpace = layout.getProgramMemorySpace();
       Attribute globalMemorySpace = layout.getGlobalMemorySpace();
       uint64_t stackAlignment = layout.getStackAlignment();
@@ -72,6 +73,9 @@ struct TestDataLayoutQuery
                                 allocaMemorySpace == Attribute()
                                     ? builder.getUI32IntegerAttr(0)
                                     : allocaMemorySpace),
+           builder.getNamedAttr("mangling_mode", manglingMode == Attribute()
+                                                     ? builder.getStringAttr("")
+                                                     : manglingMode),
            builder.getNamedAttr("program_memory_space",
                                 programMemorySpace == Attribute()
                                     ? builder.getUI32IntegerAttr(0)
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -25,6 +25,8 @@ constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
 constexpr static llvm::StringLiteral kEndiannesKeyName = "dltest.endianness";
 constexpr static llvm::StringLiteral kAllocaKeyName =
     "dltest.alloca_memory_space";
+constexpr static llvm::StringLiteral kManglingModeKeyName =
+    "dltest.mangling_mode";
 constexpr static llvm::StringLiteral kProgramKeyName =
     "dltest.program_memory_space";
 constexpr static llvm::StringLiteral kGlobalKeyName =
@@ -83,6 +85,9 @@ struct CustomDataLayoutSpec
   StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
     return Builder(context).getStringAttr(kAllocaKeyName);
   }
+  StringAttr getManglingModeIdentifier(MLIRContext *context) const {
+    return Builder(context).getStringAttr(kManglingModeKeyName);
+  }
   StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const {
     return Builder(context).getStringAttr(kProgramKeyName);
   }
@@ -474,6 +479,7 @@ module {}
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
+  EXPECT_EQ(layout.getManglingMode(), Attribute());
 }
 
 TEST(DataLayout, NullSpec) {
@@ -506,6 +512,7 @@ TEST(DataLayout, NullSpec) {
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
+  EXPECT_EQ(layout.getManglingMode(), Attribute());
 
   EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
@@ -546,6 +553,7 @@ TEST(DataLayout, EmptySpec) {
   EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
+  EXPECT_EQ(layout.getManglingMode(), Attribute());
 
   EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
@@ -567,7 +575,8 @@ TEST(DataLayout, SpecWithEntries) {
   #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>,
   #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.stack_alignment", 128 : i32>,
+  #dlti.dl_entry<"dltest.mangling_mode", "o">
 > } : () -> ()
   )MLIR";
 
@@ -604,6 +613,7 @@ TEST(DataLayout, SpecWithEntries) {
   EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3));
   EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2));
   EXPECT_EQ(layout.getStackAlignment(), 128u);
+  EXPECT_EQ(layout.getManglingMode(), Builder(&ctx).getStringAttr("o"));
 }
 
 TEST(DataLayout, SpecWithTargetSystemDescEntries) {
