[MLIR][LLVM] More on CG Profile: support null function entries

[bcardosolopes]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Bruno Cardoso Lopes (bcardosolopes)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/137269.diff

6 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td (+11-4)
• (modified) mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp (+12-5)
• (modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+8-3)
• (modified) mlir/test/Dialect/LLVMIR/module-roundtrip.mlir (+2-2)
• (modified) mlir/test/Target/LLVMIR/Import/module-flags.ll (+3-3)
• (modified) mlir/test/Target/LLVMIR/llvmir.mlir (+2-2)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index cfa6ebf3e6775..7a730be974df4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -1370,10 +1370,17 @@ def ModuleFlagCGProfileEntryAttr
           ]>]
     ```
   }];
-  let parameters = (ins "FlatSymbolRefAttr":$from,
-                        "FlatSymbolRefAttr":$to,
-                        "uint64_t":$count);
-  let assemblyFormat = "`<` struct(params) `>`";
+  let parameters = (
+    ins OptionalParameter<"std::optional<FlatSymbolRefAttr>">:$from,
+        OptionalParameter<"std::optional<FlatSymbolRefAttr>">:$to,
+        "uint64_t":$count);
+  let assemblyFormat = [{
+    `<`
+      `from` `=` ($from^) : (`null`)? `,`
+      `to` `=` ($to^) : (`null`)? `,`
+      `count` `=` $count
+    `>`
+  }];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index e4b2ebee0d9d3..e8e29aa059285 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -281,12 +281,19 @@ convertModuleFlagValue(StringRef key, ArrayAttr arrayAttr,
 
   if (key == LLVMDialect::getModuleFlagKeyCGProfileName()) {
     for (auto entry : arrayAttr.getAsRange<ModuleFlagCGProfileEntryAttr>()) {
-      llvm::Function *fromFn =
-          moduleTranslation.lookupFunction(entry.getFrom().getValue());
-      llvm::Function *toFn =
-          moduleTranslation.lookupFunction(entry.getTo().getValue());
+      llvm::Metadata *fromMetadata =
+          entry.getFrom()
+              ? llvm::ValueAsMetadata::get(moduleTranslation.lookupFunction(
+                    entry.getFrom()->getValue()))
+              : nullptr;
+      llvm::Metadata *toMetadata =
+          entry.getTo()
+              ? llvm::ValueAsMetadata::get(
+                    moduleTranslation.lookupFunction(entry.getTo()->getValue()))
+              : nullptr;
+
       llvm::Metadata *vals[] = {
-          llvm::ValueAsMetadata::get(fromFn), llvm::ValueAsMetadata::get(toFn),
+          fromMetadata, toMetadata,
           mdb.createConstant(llvm::ConstantInt::get(
               llvm::Type::getInt64Ty(context), entry.getCount()))};
       nodes.push_back(llvm::MDNode::get(context, vals));
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index d73c84af48b25..e25de9d18963b 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -521,8 +521,12 @@ void ModuleImport::addDebugIntrinsic(llvm::CallInst *intrinsic) {
 
 static Attribute convertCGProfileModuleFlagValue(ModuleOp mlirModule,
                                                  llvm::MDTuple *mdTuple) {
-  auto getFunctionSymbol = [&](const llvm::MDOperand &funcMDO) {
-    auto *f = cast<llvm::ValueAsMetadata>(funcMDO);
+  auto getFunctionSymbol =
+      [&](const llvm::MDOperand &funcMDO) -> std::optional<FlatSymbolRefAttr> {
+    auto *f = dyn_cast_or_null<llvm::ValueAsMetadata>(funcMDO);
+    // nullptr is a valid value for the function pointer.
+    if (!f)
+      return std::nullopt;
     auto *llvmFn = cast<llvm::Function>(f->getValue()->stripPointerCasts());
     return FlatSymbolRefAttr::get(mlirModule->getContext(), llvmFn->getName());
   };
@@ -570,7 +574,8 @@ LogicalResult ModuleImport::convertModuleFlagsMetadata() {
 
     if (!valAttr) {
       emitWarning(mlirModule.getLoc())
-          << "unsupported module flag value: " << diagMD(val, llvmModule.get());
+          << "unsupported module flag value for key '" << key->getString()
+          << "' : " << diagMD(val, llvmModule.get());
       continue;
     }
 
diff --git a/mlir/test/Dialect/LLVMIR/module-roundtrip.mlir b/mlir/test/Dialect/LLVMIR/module-roundtrip.mlir
index b45c61ff10b74..1450fada8a990 100644
--- a/mlir/test/Dialect/LLVMIR/module-roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/module-roundtrip.mlir
@@ -9,7 +9,7 @@ module {
                      #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">,
                      #llvm.mlir.module_flag<append, "CG Profile", [
                        #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
-                       #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+                       #llvm.cgprofile_entry<from = @from, to = null, count = 222>,
                        #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
                     ]>]
 }
@@ -23,6 +23,6 @@ module {
 // CHECK-SAME: #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">,
 // CHECK-SAME: #llvm.mlir.module_flag<append, "CG Profile", [
 // CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
-// CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+// CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = null, count = 222>,
 // CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
 // CHECK-SAME: ]>]
diff --git a/mlir/test/Target/LLVMIR/Import/module-flags.ll b/mlir/test/Target/LLVMIR/Import/module-flags.ll
index 31ab8afb7ed83..fc963f0f6fa99 100644
--- a/mlir/test/Target/LLVMIR/Import/module-flags.ll
+++ b/mlir/test/Target/LLVMIR/Import/module-flags.ll
@@ -19,7 +19,7 @@
 ; CHECK-SAME: #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">]
 
 ; // -----
-; expected-warning@-2 {{unsupported module flag value: !4 = !{!"foo", i32 1}}}
+; expected-warning@-2 {{unsupported module flag value for key 'qux' : !4 = !{!"foo", i32 1}}}
 !10 = !{ i32 1, !"foo", i32 1 }
 !11 = !{ i32 4, !"bar", i32 37 }
 !12 = !{ i32 2, !"qux", i32 42 }
@@ -36,11 +36,11 @@ declare void @to()
 !20 = !{i32 5, !"CG Profile", !21}
 !21 = distinct !{!22, !23, !24}
 !22 = !{ptr @from, ptr @to, i64 222}
-!23 = !{ptr @from, ptr @from, i64 222}
+!23 = !{ptr @from, null, i64 222}
 !24 = !{ptr @to, ptr @from, i64 222}
 
 ; CHECK: llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", [
 ; CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
-; CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+; CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = null, count = 222>,
 ; CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
 ; CHECK-SAME: ]>]
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index ffd1071e872ec..fab16d560a9e4 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2840,7 +2840,7 @@ module {
 
 llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", [
   #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
-  #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+  #llvm.cgprofile_entry<from = @from, to = null, count = 222>,
   #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
 ]>]
 llvm.func @from(i32)
@@ -2851,7 +2851,7 @@ llvm.func @to()
 // CHECK: ![[#CGPROF]] = !{i32 5, !"CG Profile", ![[#LIST:]]}
 // CHECK: ![[#LIST]] = distinct !{![[#ENTRY_A:]], ![[#ENTRY_B:]], ![[#ENTRY_C:]]}
 // CHECK: ![[#ENTRY_A]] = !{ptr @from, ptr @to, i64 222}
-// CHECK: ![[#ENTRY_B]] = !{ptr @from, ptr @from, i64 222}
+// CHECK: ![[#ENTRY_B]] = !{ptr @from, null, i64 222}
 // CHECK: ![[#ENTRY_C]] = !{ptr @to, ptr @from, i64 222}
 // CHECK: ![[#DBG]] = !{i32 2, !"Debug Info Version", i32 3}
 

[gysit]

LGTM

I would probably use OptionalParameter<"FlatSymbolRefAttr"> if possible. That way the custom parser and printer may not be necessary.