[MLIR][DLTI] Add mangling style

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;
 

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

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;

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 "

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

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)

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(

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

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

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

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"
 
 ; // -----

mlir/test/Target/LLVMIR/Import/import-failure.ll

@@ -356,6 +356,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() {

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>>,

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)

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