[llvm-extract] Delete dead CG Profile edges (#134940)

When `llvm-extract`-ing a function, and the `CG Profile` flag is present in the original module, we end up with lots of `!{null, null, i64 1234}` entries for call edges that have disappeared as result of the removed functions.

This patch fixes that by adding a pass to `llvm-extract` that finds `CG Profile` edges with one or both operands `null` and removes them. This results in a cleaner output.

Author: mtrofin

llvm/include/llvm/Transforms/IPO/StripSymbols.h

@@ -42,6 +42,10 @@ struct StripDeadDebugInfoPass : PassInfoMixin<StripDeadDebugInfoPass> {
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+struct StripDeadCGProfilePass : PassInfoMixin<StripDeadCGProfilePass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_STRIPSYMBOLS_H

llvm/lib/Passes/PassRegistry.def

@@ -155,6 +155,7 @@ MODULE_PASS("strip-dead-prototypes", StripDeadPrototypesPass())
 MODULE_PASS("strip-debug-declare", StripDebugDeclarePass())
 MODULE_PASS("strip-nondebug", StripNonDebugSymbolsPass())
 MODULE_PASS("strip-nonlinetable-debuginfo", StripNonLineTableDebugInfoPass())
+MODULE_PASS("strip-dead-cg-profile", StripDeadCGProfilePass())
 MODULE_PASS("trigger-crash-module", TriggerCrashModulePass())
 MODULE_PASS("trigger-verifier-error", TriggerVerifierErrorPass())
 MODULE_PASS("tsan-module", ModuleThreadSanitizerPass())

llvm/lib/Transforms/IPO/StripSymbols.cpp

@@ -26,10 +26,12 @@
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/TypeFinder.h"
 #include "llvm/IR/ValueSymbolTable.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO/StripSymbols.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -297,3 +299,21 @@ PreservedAnalyses StripDeadDebugInfoPass::run(Module &M,
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }
+
+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::all_of(EdgeAsNode->operands(),
+                       [](const Metadata *V) { return V != nullptr; }))
+        ValidCGEdges.push_back(Edge);
+  }
+  M.setModuleFlag(Module::Append, "CG Profile",
+                  MDTuple::getDistinct(M.getContext(), ValidCGEdges));
+  return PreservedAnalyses::none();
+}

llvm/test/Transforms/StripSymbols/strip-cg-profile.ll

@@ -0,0 +1,53 @@
+; RUN: opt -passes=strip-dead-cg-profile %s -S -o - | FileCheck %s --check-prefix=NOOP
+; RUN: llvm-extract %s -func=a -S -o - | FileCheck %s --check-prefix=EXTRACT-A
+; RUN: llvm-extract %s -func=a --func=b -S -o - | FileCheck %s --check-prefix=EXTRACT-AB
+; RUN: llvm-extract %s -func=solo -S -o - | FileCheck %s --check-prefix=NOTHING-LEFT
+
+define void @a() {
+  call void @b()
+  ret void
+}
+
+define void @b() {
+  call void @c()
+  ret void
+}
+
+define void @c() {
+  call void @d()
+  ret void
+}
+
+define void @d() {
+  ret void
+}
+
+define void @solo() {
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 5, !"CG Profile", !1}
+!1 = !{!2, !3, !4}
+!2 = !{ptr @a, ptr @b, i64 42}
+!3 = !{ptr @b, ptr @c, i64 20}
+!4 = !{ptr @c, ptr @d, i64 101}
+
+; NOOP:       !0 = !{i32 5, !"CG Profile", !1}
+; NOOP-NEXT:  !1 = distinct !{!2, !3, !4}
+; NOOP-NEXT:  !2 = !{ptr @a, ptr @b, i64 42}
+; NOOP-NEXT:  !3 = !{ptr @b, ptr @c, i64 20}
+; NOOP-NEXT:  !4 = !{ptr @c, ptr @d, i64 101}
+
+; EXTRACT-A:      !0 = !{i32 5, !"CG Profile", !1}
+; EXTRACT-A-NEXT: !1 = distinct !{!2}
+; EXTRACT-A-NEXT: !2 = !{ptr @a, ptr @b, i64 42}
+
+; EXTRACT-AB:      !0 = !{i32 5, !"CG Profile", !1}
+; EXTRACT-AB-NEXT: !1 = distinct !{!2, !3}
+; EXTRACT-AB-NEXT: !2 = !{ptr @a, ptr @b, i64 42}
+; EXTRACT-AB-NEXT: !3 = !{ptr @b, ptr @c, i64 20}
+
+; NOTHING-LEFT:      !0 = !{i32 5, !"CG Profile", !1}
+; NOTHING-LEFT-NEXT: !1 = distinct !{}

llvm/tools/llvm-extract/llvm-extract.cpp

@@ -408,6 +408,7 @@ int main(int argc, char **argv) {
     PM.addPass(GlobalDCEPass());
   PM.addPass(StripDeadDebugInfoPass());
   PM.addPass(StripDeadPrototypesPass());
+  PM.addPass(StripDeadCGProfilePass());
 
   std::error_code EC;
   ToolOutputFile Out(OutputFilename, EC, sys::fs::OF_None);