From 7bdd0a1c66707cfcab92348a8ef0d45911448788 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin@google.com>
Date: Tue, 8 Apr 2025 09:45:07 -0700
Subject: [PATCH] Strip CG Profile

---
 .../llvm/Transforms/IPO/StripSymbols.h        |  4 ++
 llvm/lib/Passes/PassRegistry.def              |  1 +
 llvm/lib/Transforms/IPO/StripSymbols.cpp      | 20 +++++++
 .../StripSymbols/strip-cg-profile.ll          | 53 +++++++++++++++++++
 llvm/tools/llvm-extract/llvm-extract.cpp      |  1 +
 5 files changed, 79 insertions(+)
 create mode 100644 llvm/test/Transforms/StripSymbols/strip-cg-profile.ll

diff --git a/llvm/include/llvm/Transforms/IPO/StripSymbols.h b/llvm/include/llvm/Transforms/IPO/StripSymbols.h
index dd76d481d668c..bd5cdde290dde 100644
--- a/llvm/include/llvm/Transforms/IPO/StripSymbols.h
+++ b/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
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 94f5bf2c7c6c5..d5d1b2173da69 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/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())
diff --git a/llvm/lib/Transforms/IPO/StripSymbols.cpp b/llvm/lib/Transforms/IPO/StripSymbols.cpp
index e7b70308ba7aa..647caaa123cf8 100644
--- a/llvm/lib/Transforms/IPO/StripSymbols.cpp
+++ b/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();
+}
\ No newline at end of file
diff --git a/llvm/test/Transforms/StripSymbols/strip-cg-profile.ll b/llvm/test/Transforms/StripSymbols/strip-cg-profile.ll
new file mode 100644
index 0000000000000..89176348eef9d
--- /dev/null
+++ b/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 !{}
\ No newline at end of file
diff --git a/llvm/tools/llvm-extract/llvm-extract.cpp b/llvm/tools/llvm-extract/llvm-extract.cpp
index 169cd0c2e4cbf..648060acb392c 100644
--- a/llvm/tools/llvm-extract/llvm-extract.cpp
+++ b/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);