[ctxprof] Only prune the profile in modules containing only context trees

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -33,6 +33,11 @@ class PGOContextualProfile {
     FunctionInfo(StringRef Name) : Name(Name) {}
   };
   PGOCtxProfile Profiles;
+
+  // True if this module is a post-thinlto module containing just functions
+  // participating in one or more contextual profiles.
+  bool IsInSpecializedModule = false;
+
   // For the GUIDs in this module, associate metadata about each function which
   // we'll need when we maintain the profiles during IPO transformations.
   std::map<GlobalValue::GUID, FunctionInfo> FuncInfo;
@@ -56,6 +61,8 @@ class PGOContextualProfile {
 
   const PGOCtxProfile &profiles() const { return Profiles; }
 
+  bool isInSpecializedModule() const { return IsInSpecializedModule; }
+
   bool isFunctionKnown(const Function &F) const {
     return getDefinedFunctionGUID(F) != 0;
   }

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/CtxProfAnalysis.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Analysis.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -20,6 +21,7 @@
 #include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 
 #define DEBUG_TYPE "ctx_prof"
 
@@ -95,26 +97,42 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
     return {};
   }
 
-  DenseSet<GlobalValue::GUID> ProfileRootsInModule;
-  for (const auto &F : M)
-    if (!F.isDeclaration())
-      if (auto GUID = AssignGUIDPass::getGUID(F);
-          MaybeProfiles->Contexts.find(GUID) != MaybeProfiles->Contexts.end())
-        ProfileRootsInModule.insert(GUID);
-
-  // Trim first the roots that aren't in this module.
-  for (auto &[RootGuid, _] :
-       llvm::make_early_inc_range(MaybeProfiles->Contexts))
-    if (!ProfileRootsInModule.contains(RootGuid))
-      MaybeProfiles->Contexts.erase(RootGuid);
-  // If none of the roots are in the module, we have no profile (for this
-  // module)
-  if (MaybeProfiles->Contexts.empty())
-    return {};
-
-  // OK, so we have a valid profile and it's applicable to roots in this module.
+  // FIXME: We should drive this from ThinLTO, but for the time being, use the
+  // module name as indicator.
+  // We want to *only* keep the contextual profiles in modules that capture
+  // context trees. That allows us to compute specific PSIs, for example.
+  auto DetermineRootsInModule = [&M]() -> const DenseSet<GlobalValue::GUID> {
+    DenseSet<GlobalValue::GUID> ProfileRootsInModule;
+    auto ModName = M.getName();
+    auto Filename = sys::path::filename(ModName);
+    // Drop the file extension.
+    Filename = Filename.substr(0, Filename.find_last_of('.'));
+    // See if it parses
+    APInt Guid;
+    // getAsInteger returns true if there are more chars to read other than the
+    // integer. So the "false" test is what we want.
+    if (!Filename.getAsInteger(0, Guid))
+      ProfileRootsInModule.insert(Guid.getZExtValue());
+    return ProfileRootsInModule;
+  };
+  const auto ProfileRootsInModule = DetermineRootsInModule();
   PGOContextualProfile Result;
 
+  // the logic from here on allows for modules that contain - by design - more
+  // than one root. We currently don't support that, because the determination
+  // happens based on the module name matching the root guid, but the logic can
+  // avoid assuming that.
+  if (!ProfileRootsInModule.empty()) {
+    Result.IsInSpecializedModule = true;
+    // Trim first the roots that aren't in this module.
+    for (auto &[RootGuid, _] :
+         llvm::make_early_inc_range(MaybeProfiles->Contexts))
+      if (!ProfileRootsInModule.contains(RootGuid))
+        MaybeProfiles->Contexts.erase(RootGuid);
+    // we can also drop the flat profiles
+    MaybeProfiles->FlatProfiles.clear();
+  }
+
   for (const auto &F : M) {
     if (F.isDeclaration())
       continue;

llvm/lib/Transforms/IPO/ElimAvailExtern.cpp

@@ -134,7 +134,7 @@ EliminateAvailableExternallyPass::run(Module &M, ModuleAnalysisManager &MAM) {
   // for this contextual information. Eliding it in favor of the original would
   // undo these optimizations.
   if (!eliminateAvailableExternally(
-          M, /*Convert=*/(CtxProf && !CtxProf->contexts().empty())))
+          M, /*Convert=*/(CtxProf && CtxProf->isInSpecializedModule())))
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }

llvm/test/Analysis/CtxProfAnalysis/load.ll

@@ -46,6 +46,9 @@ Function Info:
 Current Profile:
 
 Contexts:
+  - Guid:            12341
+    TotalRootEntryCount: 90
+    Counters:        [ 9 ]
   - Guid:            11872291593386833696
     TotalRootEntryCount: 4
     Counters:        [ 1 ]
@@ -57,6 +60,7 @@ Contexts:
     Counters:        [ 5 ]
 
 Flat Profile:
+12341 : 9 
 728453322856651412 : 6 7 
 11872291593386833696 : 1 
 12074870348631550642 : 5 

llvm/test/Analysis/CtxProfAnalysis/pruning.ll

@@ -0,0 +1,69 @@
+; REQUIRES: x86_64-linux
+;
+; Check that we don't prune the contextual profile, unless the module name
+; matches the guid of the root.
+;
+; RUN: rm -rf %t
+; RUN: split-file %s %t
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
+;
+; RUN: cp %t/example.ll %t/1234.ll
+; RUN: cp %t/example.ll %t/0x4d2.ll
+;
+; RUN: opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -ctx-profile-printer-level=everything \
+; RUN:   %t/example.ll -S 2>&1 | FileCheck %s
+
+; RUN: opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -ctx-profile-printer-level=everything \
+; RUN:   %t/not-matching.ll -S 2>&1 | FileCheck %s
+
+; RUN: opt -passes='require<ctx-prof-analysis>,print<ctx-prof-analysis>' \
+; RUN:   -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   -ctx-profile-printer-level=everything \
+; RUN:   %t/0x4d2.ll -S 2>&1 | FileCheck %s --check-prefix=PRUNED
+
+; CHECK:    Contexts:
+; CHECK:      - Guid:            1234
+; CHECK:      - Guid:            5678
+; CHECK:    FlatProfiles:
+; PRUNED-NOT: - Guid:            5678
+; PRUNED-NOT: FlatProfiles
+;
+; pick a large GUID that would be negative, if signed, to test a few ways the
+; file name may be formatted.
+;--- profile.yaml
+Contexts:
+  - Guid: 1234
+    TotalRootEntryCount: 24
+    Counters: [9]
+    Callsites:  -
+                  - Guid: 1000
+                    Counters: [6, 7]
+
+  - Guid: 5678
+    TotalRootEntryCount: 24
+    Counters: [9]
+    Callsites:  -
+                  - Guid: 1000
+                    Counters: [6, 7]
+FlatProfiles:
+  - Guid: 777
+    Counters: [2]                    
+;--- example.ll
+define void @an_entrypoint(i32 %a) !guid !0 {
+  ret void
+}
+
+attributes #0 = { noinline }
+!0 = !{ i64  1234 }
+
+;--- not-matching.ll
+define void @an_entrypoint(i32 %a) !guid !0 {
+  ret void
+}
+
+attributes #0 = { noinline }
+!0 = !{ i64  1000 }

llvm/test/Transforms/EliminateAvailableExternally/transform-to-local.ll

@@ -1,16 +1,26 @@
 ; REQUIRES: asserts
-; RUN: opt -passes=elim-avail-extern -avail-extern-to-local -stats -S 2>&1 < %s | FileCheck %s
+; RUN: rm -rf %t
+; RUN: mkdir %t
+; RUN: cp %s %t/1234.ll
+;
+; default behavior
+; RUN: opt -passes=elim-avail-extern -stats -S 2>&1 %s | FileCheck %s --check-prefix=NOOP
+;
+; check the -avail-extern-to-local flag works as intended
+; RUN: opt -passes=elim-avail-extern -avail-extern-to-local -stats -S 2>&1 %t/1234.ll | FileCheck %s
+; RUN: opt -passes=elim-avail-extern -avail-extern-to-local -stats -S 2>&1 %s | FileCheck %s
 ;
 ; RUN: echo '{"Contexts": [{"Guid":1234, "TotalRootEntryCount": 5, "Counters": [1]}]}' | llvm-ctxprof-util fromYAML --input=- --output=%t_profile.ctxprofdata
 ;
 ; Because we pass a contextual profile with a root defined in this module, we expect the outcome to be the same as-if
 ; we passed -avail-extern-to-local, i.e. available_externally don't get elided and instead get converted to local linkage
-; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile.ctxprofdata -stats -S 2>&1 < %s | FileCheck %s
+;
+; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile.ctxprofdata -stats -S 2>&1 %t/1234.ll | FileCheck %s
 
 ; If the profile doesn't apply to this module, available_externally won't get converted to internal linkage, and will be
 ; removed instead.
 ; RUN: echo '{"Contexts": [{"Guid":5678, "TotalRootEntryCount": 3, "Counters": [1]}]}' | llvm-ctxprof-util fromYAML --input=- --output=%t_profile_bad.ctxprofdata
-; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile_bad.ctxprofdata -stats -S 2>&1 < %s | FileCheck %s --check-prefix=NOOP
+; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile_bad.ctxprofdata -stats -S %s 2>&1 | FileCheck %s --check-prefix=NOOP
 
 declare void @call_out(ptr %fct)