[mlir][LLVMIR] Handle missing functions in CGProfile module flags

[Men-cotton]

Fixes: #160717

Now build/bin/mlir-translate -mlir-to-llvmir mlir/test/Dialect/LLVMIR/module-roundtrip.mlir passes.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Men-cotton (Men-cotton)

Changes

Fixes: #160717

Now build/bin/mlir-translate -mlir-to-llvmir mlir/test/Dialect/LLVMIR/module-roundtrip.mlir passes.

---

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

2 Files Affected:

• (modified) mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp (+11-10)
• (added) mlir/test/Dialect/LLVMIR/invalid-cg-profile.mlir (+26)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 7d3486acaf82a..f11b9a6c61690 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -243,16 +243,17 @@ convertModuleFlagValue(StringRef key, ArrayAttr arrayAttr,
 
   if (key == LLVMDialect::getModuleFlagKeyCGProfileName()) {
     for (auto entry : arrayAttr.getAsRange<ModuleFlagCGProfileEntryAttr>()) {
-      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;
+      const auto getFuncMetadata =
+          [&](FlatSymbolRefAttr sym) -> llvm::Metadata * {
+        if (!sym)
+          return nullptr;
+        if (llvm::Function *fn =
+                moduleTranslation.lookupFunction(sym.getValue()))
+          return llvm::ValueAsMetadata::get(fn);
+        return nullptr;
+      };
+      llvm::Metadata *const fromMetadata = getFuncMetadata(entry.getFrom());
+      llvm::Metadata *const toMetadata = getFuncMetadata(entry.getTo());
 
       llvm::Metadata *vals[] = {
           fromMetadata, toMetadata,
diff --git a/mlir/test/Dialect/LLVMIR/invalid-cg-profile.mlir b/mlir/test/Dialect/LLVMIR/invalid-cg-profile.mlir
new file mode 100644
index 0000000000000..d7a0f386303e0
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/invalid-cg-profile.mlir
@@ -0,0 +1,26 @@
+// RUN: mlir-translate %s -mlir-to-llvmir | FileCheck %s
+// CHECK: !llvm.module.flags
+
+module {
+  llvm.module_flags [#llvm.mlir.module_flag<error, "wchar_size", 4 : i32>,
+                     #llvm.mlir.module_flag<min, "PIC Level", 2 : i32>,
+                     #llvm.mlir.module_flag<max, "PIE Level", 2 : i32>,
+                     #llvm.mlir.module_flag<max, "uwtable", 2 : i32>,
+                     #llvm.mlir.module_flag<max, "frame-pointer", 1 : i32>,
+                     #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, count = 222>,
+                       #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+                    ]>,
+                    #llvm.mlir.module_flag<error, "ProfileSummary",
+                       #llvm.profile_summary<format = InstrProf, total_count = 263646, max_count = 86427,
+                         max_internal_count = 86427, max_function_count = 4691,
+                         num_counts = 3712, num_functions = 796,
+                         is_partial_profile = 0,
+                         partial_profile_ratio = 0.000000e+00 : f64,
+                         detailed_summary =
+                           <cut_off = 10000, min_count = 86427, num_counts = 1>,
+                           <cut_off = 100000, min_count = 86427, num_counts = 1>
+                    >>]
+}

[Men-cotton]

CC: @gysit, @Dinistro

[Dinistro]

Do you know if it's legal that the cgprofile entries have missing functions?

[Dinistro]

Do you know if it's legal that the cgprofile entries have missing functions?

@Men-cotton can you give some context on this? I did not perform a deep search but so far did not encounter any documentation that states anything about this.

[Men-cotton]

Do you know if it's legal that the cgprofile entries have missing functions?

I think it's legal.

llvm-project/llvm/lib/IR/Verifier.cpp

Lines 2033 to 2048 in cca66a2

	void Verifier::visitModuleFlagCGProfileEntry(const MDOperand &MDO) {
	auto CheckFunction = [&](const MDOperand &FuncMDO) {
	if (!FuncMDO)
	return;
	auto F = dyn_cast<ValueAsMetadata>(FuncMDO);
	Check(F && isa<Function>(F->getValue()->stripPointerCasts()),
	"expected a Function or null", FuncMDO);
	};
	auto Node = dyn_cast_or_null<MDNode>(MDO);
	Check(Node && Node->getNumOperands() == 3, "expected a MDNode triple", MDO);
	CheckFunction(Node->getOperand(0));
	CheckFunction(Node->getOperand(1));
	auto Count = dyn_cast_or_null<ConstantAsMetadata>(Node->getOperand(2));
	Check(Count && Count->getType()->isIntegerTy(),
	"expected an integer constant", Node->getOperand(2));
	}

llvm-project/llvm/lib/Target/TargetLoweringObjectFile.cpp

Lines 177 to 192 in cca66a2

	for (const auto &Edge : CGProfile->operands()) {
	MDNode *E = cast<MDNode>(Edge);
	const MCSymbol *From = GetSym(E->getOperand(0));
	const MCSymbol *To = GetSym(E->getOperand(1));
	// Skip null functions. This can happen if functions are dead stripped after
	// the CGProfile pass has been run.
	if (!From || !To)
	continue;
	uint64_t Count = cast<ConstantAsMetadata>(E->getOperand(2))
	->getValue()
	->getUniqueInteger()
	.getZExtValue();
	Streamer.emitCGProfileEntry(MCSymbolRefExpr::create(From, C),
	MCSymbolRefExpr::create(To, C), Count);
	}
	}

[Men-cotton]

One more relevant spot: StripDeadCGProfilePass drops null edges before codegen. It rebuilds the CG Profile flag keeping only edges whose operands contain no nullptr:

llvm-project/llvm/lib/Transforms/IPO/StripSymbols.cpp

Lines 303 to 318 in cca66a2

	PreservedAnalyses StripDeadCGProfilePass::run(Module &M,
	ModuleAnalysisManager &AM) {
	auto *CGProf = dyn_cast_or_null<MDTuple>(M.getModuleFlag("CG Profile"));
	if (!CGProf)
	return PreservedAnalyses::all();
	SmallVector<Metadata *, 16> ValidCGEdges;
	for (Metadata *Edge : CGProf->operands()) {
	if (auto *EdgeAsNode = dyn_cast_or_null<MDNode>(Edge))
	if (!llvm::is_contained(EdgeAsNode->operands(), nullptr))
	ValidCGEdges.push_back(Edge);
	}
	M.setModuleFlag(Module::Append, "CG Profile",
	MDTuple::getDistinct(M.getContext(), ValidCGEdges));
	return PreservedAnalyses::none();
	}

So verifier allows nulls, object emission skips them, and this pass proactively strips them.

[gysit]

LGTM modulo the @Dinistro 's comment!

Thanks for the pointers!

[Dinistro]

LGTM, thanks for the clarifications and addressing the comments.