From 6cb2d40387f613c72ec8f66e47a5eb0f7e4b0f1f Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin@google.com>
Date: Fri, 6 Sep 2024 13:44:05 -0700
Subject: [PATCH] [ctx_prof] Handle case when no root is in this Module.
 (#107463)

If none of the functions in this `Module` are roots in the contextual profile, we can't use it and should just return the `{}` case.
---
 llvm/lib/Analysis/CtxProfAnalysis.cpp         | 21 +++++--
 .../CtxProfAnalysis/load-unapplicable.ll      | 63 +++++++++++++++++++
 2 files changed, 80 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Analysis/CtxProfAnalysis/load-unapplicable.ll

diff --git a/llvm/lib/Analysis/CtxProfAnalysis.cpp b/llvm/lib/Analysis/CtxProfAnalysis.cpp
index 457a4dcc796847..560d9c742d5e7d 100644
--- a/llvm/lib/Analysis/CtxProfAnalysis.cpp
+++ b/llvm/lib/Analysis/CtxProfAnalysis.cpp
@@ -132,6 +132,23 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
     return {};
   }
 
+  DenseSet<GlobalValue::GUID> ProfileRootsInModule;
+  for (const auto &F : M)
+    if (!F.isDeclaration())
+      if (auto GUID = AssignGUIDPass::getGUID(F);
+          MaybeCtx->find(GUID) != MaybeCtx->end())
+        ProfileRootsInModule.insert(GUID);
+
+  // Trim first the roots that aren't in this module.
+  for (auto &[RootGuid, _] : llvm::make_early_inc_range(*MaybeCtx))
+    if (!ProfileRootsInModule.contains(RootGuid))
+      MaybeCtx->erase(RootGuid);
+  // If none of the roots are in the module, we have no profile (for this
+  // module)
+  if (MaybeCtx->empty())
+    return {};
+
+  // OK, so we have a valid profile and it's applicable to roots in this module.
   PGOContextualProfile Result;
 
   for (const auto &F : M) {
@@ -166,10 +183,6 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
   }
   // If we made it this far, the Result is valid - which we mark by setting
   // .Profiles.
-  // Trim first the roots that aren't in this module.
-  for (auto &[RootGuid, _] : llvm::make_early_inc_range(*MaybeCtx))
-    if (!Result.FuncInfo.contains(RootGuid))
-      MaybeCtx->erase(RootGuid);
   Result.Profiles = std::move(*MaybeCtx);
   return Result;
 }
diff --git a/llvm/test/Analysis/CtxProfAnalysis/load-unapplicable.ll b/llvm/test/Analysis/CtxProfAnalysis/load-unapplicable.ll
new file mode 100644
index 00000000000000..ac5e7930d3f532
--- /dev/null
+++ b/llvm/test/Analysis/CtxProfAnalysis/load-unapplicable.ll
@@ -0,0 +1,63 @@
+; REQUIRES: x86_64-linux
+;
+; Check that, if none of the roots in the profile are defined in the module, the
+; profile is treated as empty (i.e. "none provided")
+;
+; RUN: rm -rf %t
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' -ctx-profile-printer-level=everything \
+; RUN:   %t/example.ll -S 2>&1 | FileCheck %s
+
+; CHECK: No contextual profile was provided
+;
+; This is the reference profile, laid out in the format the json formatter will
+; output it from opt.
+;--- profile.json
+[
+  {
+    "Counters": [
+      9
+    ],
+    "Guid": 12341
+  },
+  {
+    "Counters": [
+      5
+    ],
+    "Guid": 12074870348631550642
+  },
+  {
+    "Callsites": [
+      [
+        {
+          "Counters": [
+            6,
+            7
+          ],
+          "Guid": 728453322856651412
+        }
+      ]
+    ],
+    "Counters": [
+      1
+    ],
+    "Guid": 11872291593386833696
+  }
+]
+;--- example.ll
+declare void @bar()
+
+define void @an_entrypoint(i32 %a) !guid !0 {
+  %t = icmp eq i32 %a, 0
+  br i1 %t, label %yes, label %no
+
+yes:
+  call void @bar()
+  ret void
+no:
+  ret void
+}
+
+attributes #0 = { noinline }
+!0 = !{ i64 1000 }
\ No newline at end of file