[MLIR][LLVM] More on CG Profile: support null function entries (llvm#137269)

Author: bcardosolopes

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -1370,9 +1370,11 @@ def ModuleFlagCGProfileEntryAttr
           ]>]
     ```
   }];
-  let parameters = (ins "FlatSymbolRefAttr":$from,
-                        "FlatSymbolRefAttr":$to,
-                        "uint64_t":$count);
+  let parameters = (
+    ins OptionalParameter<"FlatSymbolRefAttr">:$from,
+        OptionalParameter<"FlatSymbolRefAttr">:$to,
+        "uint64_t":$count);
+
   let assemblyFormat = "`<` struct(params) `>`";
 }
 

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

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -521,10 +521,14 @@ 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 getLLVMFunction =
+      [&](const llvm::MDOperand &funcMDO) -> llvm::Function * {
+    auto *f = cast_or_null<llvm::ValueAsMetadata>(funcMDO);
+    // nullptr is a valid value for the function pointer.
+    if (!f)
+      return nullptr;
     auto *llvmFn = cast<llvm::Function>(f->getValue()->stripPointerCasts());
-    return FlatSymbolRefAttr::get(mlirModule->getContext(), llvmFn->getName());
+    return llvmFn;
   };
 
   // Each tuple element becomes one ModuleFlagCGProfileEntryAttr.
@@ -535,9 +539,17 @@ static Attribute convertCGProfileModuleFlagValue(ModuleOp mlirModule,
     llvm::Constant *llvmConstant =
         cast<llvm::ConstantAsMetadata>(cgEntry->getOperand(2))->getValue();
     uint64_t count = cast<llvm::ConstantInt>(llvmConstant)->getZExtValue();
+    auto *fromFn = getLLVMFunction(cgEntry->getOperand(0));
+    auto *toFn = getLLVMFunction(cgEntry->getOperand(1));
+    // FlatSymbolRefAttr::get(mlirModule->getContext(), llvmFn->getName());
     cgProfile.push_back(ModuleFlagCGProfileEntryAttr::get(
-        mlirModule->getContext(), getFunctionSymbol(cgEntry->getOperand(0)),
-        getFunctionSymbol(cgEntry->getOperand(1)), count));
+        mlirModule->getContext(),
+        fromFn ? FlatSymbolRefAttr::get(mlirModule->getContext(),
+                                        fromFn->getName())
+               : nullptr,
+        toFn ? FlatSymbolRefAttr::get(mlirModule->getContext(), toFn->getName())
+             : nullptr,
+        count));
   }
   return ArrayAttr::get(mlirModule->getContext(), cgProfile);
 }
@@ -570,7 +582,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;
     }
 

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, 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, count = 222>,
 // CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
 // CHECK-SAME: ]>]

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, count = 222>,
 ; CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
 ; CHECK-SAME: ]>]

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