[ctxprof][nfc] Prepare CtxProfAnalysis for flat profiles

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -21,12 +21,6 @@ namespace llvm {
 
 class CtxProfAnalysis;
 
-// Setting initial capacity to 1 because all contexts must have at least 1
-// counter, and then, because all contexts belonging to a function have the same
-// size, there'll be at most one other heap allocation.
-using CtxProfFlatProfile =
-    std::map<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
-
 /// The instrumented contextual profile, produced by the CtxProfAnalysis.
 class PGOContextualProfile {
   friend class CtxProfAnalysis;
@@ -38,7 +32,7 @@ class PGOContextualProfile {
     PGOCtxProfContext Index;
     FunctionInfo(StringRef Name) : Name(Name) {}
   };
-  std::optional<PGOCtxProfContext::CallTargetMapTy> Profiles;
+  PGOCtxProfile Profiles;
   // 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,10 +50,8 @@ class PGOContextualProfile {
   PGOContextualProfile(const PGOContextualProfile &) = delete;
   PGOContextualProfile(PGOContextualProfile &&) = default;
 
-  operator bool() const { return Profiles.has_value(); }
-
-  const PGOCtxProfContext::CallTargetMapTy &profiles() const {
-    return *Profiles;
+  const CtxProfContextualProfiles &contexts() const {
+    return Profiles.Contexts;
   }
 
   bool isFunctionKnown(const Function &F) const {

llvm/include/llvm/ProfileData/PGOCtxProfReader.h

@@ -164,6 +164,23 @@ class PGOCtxProfContext final : public internal::IndexNode {
   }
 };
 
+// Setting initial capacity to 1 because all contexts must have at least 1
+// counter, and then, because all contexts belonging to a function have the same
+// size, there'll be at most one other heap allocation.
+using CtxProfFlatProfile =
+    std::map<GlobalValue::GUID, SmallVector<uint64_t, 1>>;
+
+using CtxProfContextualProfiles =
+    std::map<GlobalValue::GUID, PGOCtxProfContext>;
+struct PGOCtxProfile {
+  CtxProfContextualProfiles Contexts;
+
+  PGOCtxProfile() = default;
+  PGOCtxProfile(const PGOCtxProfile &) = delete;
+  PGOCtxProfile(PGOCtxProfile &&) = default;
+  PGOCtxProfile &operator=(PGOCtxProfile &&) = default;
+};
+
 class PGOCtxProfileReader final {
   StringRef Magic;
   BitstreamCursor Cursor;
@@ -181,7 +198,7 @@ class PGOCtxProfileReader final {
       : Magic(Buffer.substr(0, PGOCtxProfileWriter::ContainerMagic.size())),
         Cursor(Buffer.substr(PGOCtxProfileWriter::ContainerMagic.size())) {}
 
-  Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>> loadContexts();
+  Expected<PGOCtxProfile> loadProfiles();
 };
 
 void convertCtxProfToYaml(raw_ostream &OS,

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -88,27 +88,28 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
     return {};
   }
   PGOCtxProfileReader Reader(MB.get()->getBuffer());
-  auto MaybeCtx = Reader.loadContexts();
-  if (!MaybeCtx) {
+  auto MaybeProfiles = Reader.loadProfiles();
+  if (!MaybeProfiles) {
     M.getContext().emitError("contextual profile file is invalid: " +
-                             toString(MaybeCtx.takeError()));
+                             toString(MaybeProfiles.takeError()));
     return {};
   }
 
   DenseSet<GlobalValue::GUID> ProfileRootsInModule;
   for (const auto &F : M)
     if (!F.isDeclaration())
       if (auto GUID = AssignGUIDPass::getGUID(F);
-          MaybeCtx->find(GUID) != MaybeCtx->end())
+          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(*MaybeCtx))
+  for (auto &[RootGuid, _] :
+       llvm::make_early_inc_range(MaybeProfiles->Contexts))
     if (!ProfileRootsInModule.contains(RootGuid))
-      MaybeCtx->erase(RootGuid);
+      MaybeProfiles->Contexts.erase(RootGuid);
   // If none of the roots are in the module, we have no profile (for this
   // module)
-  if (MaybeCtx->empty())
+  if (MaybeProfiles->Contexts.empty())
     return {};
 
   // OK, so we have a valid profile and it's applicable to roots in this module.
@@ -146,7 +147,7 @@ PGOContextualProfile CtxProfAnalysis::run(Module &M,
   }
   // If we made it this far, the Result is valid - which we mark by setting
   // .Profiles.
-  Result.Profiles = std::move(*MaybeCtx);
+  Result.Profiles = std::move(*MaybeProfiles);
   Result.initIndex();
   return Result;
 }
@@ -164,7 +165,7 @@ CtxProfAnalysisPrinterPass::CtxProfAnalysisPrinterPass(raw_ostream &OS)
 PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
                                                   ModuleAnalysisManager &MAM) {
   CtxProfAnalysis::Result &C = MAM.getResult<CtxProfAnalysis>(M);
-  if (!C) {
+  if (C.contexts().empty()) {
     OS << "No contextual profile was provided.\n";
     return PreservedAnalyses::all();
   }
@@ -179,7 +180,7 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
 
   if (Mode == PrintMode::Everything)
     OS << "\nCurrent Profile:\n";
-  convertCtxProfToYaml(OS, C.profiles());
+  convertCtxProfToYaml(OS, C.contexts());
   OS << "\n";
   if (Mode == PrintMode::YAML)
     return PreservedAnalyses::all();
@@ -245,7 +246,7 @@ void PGOContextualProfile::initIndex() {
   for (auto &[Guid, FI] : FuncInfo)
     InsertionPoints[Guid] = &FI.Index;
   preorderVisit<PGOCtxProfContext::CallTargetMapTy, PGOCtxProfContext>(
-      *Profiles, [&](PGOCtxProfContext &Ctx) {
+      Profiles.Contexts, [&](PGOCtxProfContext &Ctx) {
         auto InsertIt = InsertionPoints.find(Ctx.guid());
         if (InsertIt == InsertionPoints.end())
           return;
@@ -270,7 +271,7 @@ void PGOContextualProfile::update(Visitor V, const Function &F) {
 void PGOContextualProfile::visit(ConstVisitor V, const Function *F) const {
   if (!F)
     return preorderVisit<const PGOCtxProfContext::CallTargetMapTy,
-                         const PGOCtxProfContext>(*Profiles, V);
+                         const PGOCtxProfContext>(Profiles.Contexts, V);
   assert(isFunctionKnown(*F));
   GlobalValue::GUID G = getDefinedFunctionGUID(*F);
   for (const auto *Node = FuncInfo.find(G)->second.Index.Next; Node;
@@ -279,11 +280,10 @@ void PGOContextualProfile::visit(ConstVisitor V, const Function *F) const {
 }
 
 const CtxProfFlatProfile PGOContextualProfile::flatten() const {
-  assert(Profiles.has_value());
   CtxProfFlatProfile Flat;
   preorderVisit<const PGOCtxProfContext::CallTargetMapTy,
                 const PGOCtxProfContext>(
-      *Profiles, [&](const PGOCtxProfContext &Ctx) {
+      Profiles.Contexts, [&](const PGOCtxProfContext &Ctx) {
         auto [It, Ins] = Flat.insert({Ctx.guid(), {}});
         if (Ins) {
           llvm::append_range(It->second, Ctx.counters());

llvm/lib/ProfileData/PGOCtxProfReader.cpp

@@ -168,14 +168,13 @@ Error PGOCtxProfileReader::readMetadata() {
   return Error::success();
 }
 
-Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>>
-PGOCtxProfileReader::loadContexts() {
-  std::map<GlobalValue::GUID, PGOCtxProfContext> Ret;
+Expected<PGOCtxProfile> PGOCtxProfileReader::loadProfiles() {
+  PGOCtxProfile Ret;
   RET_ON_ERR(readMetadata());
   while (canReadContext()) {
     EXPECT_OR_RET(E, readContext(false));
     auto Key = E->second.guid();
-    if (!Ret.insert({Key, std::move(E->second)}).second)
+    if (!Ret.Contexts.insert({Key, std::move(E->second)}).second)
       return wrongValue("Duplicate roots");
   }
   return std::move(Ret);

llvm/lib/Transforms/IPO/ElimAvailExtern.cpp

@@ -133,7 +133,8 @@ EliminateAvailableExternallyPass::run(Module &M, ModuleAnalysisManager &MAM) {
   // that's imported, its optimizations will, thus, differ, and be specialized
   // for this contextual information. Eliding it in favor of the original would
   // undo these optimizations.
-  if (!eliminateAvailableExternally(M, /*Convert=*/(CtxProf && !!(*CtxProf))))
+  if (!eliminateAvailableExternally(
+          M, /*Convert=*/(CtxProf && !CtxProf->contexts().empty())))
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }

llvm/lib/Transforms/IPO/FunctionImport.cpp

@@ -724,12 +724,12 @@ class WorkloadImportsManager : public ModuleImportsManager {
     auto Buffer = std::move(BufferOrErr.get());
 
     PGOCtxProfileReader Reader(Buffer->getBuffer());
-    auto Ctx = Reader.loadContexts();
+    auto Ctx = Reader.loadProfiles();
     if (!Ctx) {
       report_fatal_error("Failed to parse contextual profiles");
       return;
     }
-    const auto &CtxMap = *Ctx;
+    const auto &CtxMap = Ctx->Contexts;
     SetVector<GlobalValue::GUID> ContainedGUIDs;
     for (const auto &[RootGuid, Root] : CtxMap) {
       // Avoid ContainedGUIDs to get in/out of scope. Reuse its memory for

llvm/lib/Transforms/IPO/ModuleInliner.cpp

@@ -171,7 +171,7 @@ PreservedAnalyses ModuleInlinerPass::run(Module &M,
                      << setIsVerbose();
             });
           }
-        } else if (CtxProfPromoteAlwaysInline && CtxProf &&
+        } else if (CtxProfPromoteAlwaysInline && !CtxProf.contexts().empty() &&
                    CB->isIndirectCall()) {
           CtxProfAnalysis::collectIndirectCallPromotionList(*CB, CtxProf,
                                                             ICPCandidates);
@@ -260,7 +260,7 @@ PreservedAnalyses ModuleInlinerPass::run(Module &M,
           // iteration because the next iteration may not happen and we may
           // miss inlining it.
           // FIXME: enable for ctxprof.
-          if (!CtxProf)
+          if (CtxProf.contexts().empty())
             if (tryPromoteCall(*ICB))
               NewCallee = ICB->getCalledFunction();
         }

llvm/lib/Transforms/Instrumentation/PGOCtxProfFlattening.cpp

@@ -435,7 +435,7 @@ PreservedAnalyses PGOCtxProfFlatteningPass::run(Module &M,
       removeInstrumentation(F);
   });
   auto &CtxProf = MAM.getResult<CtxProfAnalysis>(M);
-  if (!CtxProf)
+  if (CtxProf.contexts().empty())
     return PreservedAnalyses::none();
 
   const auto FlattenedProfile = CtxProf.flatten();

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -2357,7 +2357,7 @@ llvm::InlineResult llvm::InlineFunction(CallBase &CB, InlineFunctionInfo &IFI,
                                         AAResults *CalleeAAR,
                                         bool InsertLifetime,
                                         Function *ForwardVarArgsTo) {
-  if (!CtxProf)
+  if (CtxProf.contexts().empty())
     return InlineFunction(CB, IFI, MergeAttributes, CalleeAAR, InsertLifetime,
                           ForwardVarArgsTo);
 

llvm/tools/llvm-ctxprof-util/llvm-ctxprof-util.cpp

@@ -74,10 +74,10 @@ Error convertToYaml() {
   if (EC)
     return createStringError(EC, "failed to open output");
   PGOCtxProfileReader Reader(BufOrError.get()->getBuffer());
-  auto Prof = Reader.loadContexts();
+  auto Prof = Reader.loadProfiles();
   if (!Prof)
     return Prof.takeError();
-  llvm::convertCtxProfToYaml(Out, *Prof);
+  llvm::convertCtxProfToYaml(Out, Prof->Contexts);
   Out << "\n";
   return Error::success();
 }

llvm/unittests/ProfileData/PGOCtxProfReaderWriterTest.cpp

@@ -121,9 +121,9 @@ TEST_F(PGOCtxProfRWTest, RoundTrip) {
     EXPECT_TRUE(AnalyzerDump.find("<CalleeIndex codeid") != std::string::npos);
 
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     ASSERT_TRUE(!!Expected);
-    auto &Ctxes = *Expected;
+    auto &Ctxes = Expected->Contexts;
     EXPECT_EQ(Ctxes.size(), roots().size());
     EXPECT_EQ(Ctxes.size(), 2U);
     for (auto &[G, R] : roots())
@@ -157,22 +157,22 @@ TEST_F(PGOCtxProfRWTest, InvalidCounters) {
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
   }
 }
 
 TEST_F(PGOCtxProfRWTest, Empty) {
   PGOCtxProfileReader Reader("");
-  auto Expected = Reader.loadContexts();
+  auto Expected = Reader.loadProfiles();
   EXPECT_FALSE(Expected);
   consumeError(Expected.takeError());
 }
 
 TEST_F(PGOCtxProfRWTest, Invalid) {
   PGOCtxProfileReader Reader("Surely this is not valid");
-  auto Expected = Reader.loadContexts();
+  auto Expected = Reader.loadProfiles();
   EXPECT_FALSE(Expected);
   consumeError(Expected.takeError());
 }
@@ -194,9 +194,9 @@ TEST_F(PGOCtxProfRWTest, ValidButEmpty) {
     ASSERT_NE(*MB, nullptr);
 
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     EXPECT_TRUE(!!Expected);
-    EXPECT_TRUE(Expected->empty());
+    EXPECT_TRUE(Expected->Contexts.empty());
   }
 }
 
@@ -216,7 +216,7 @@ TEST_F(PGOCtxProfRWTest, WrongVersion) {
     ASSERT_NE(*MB, nullptr);
 
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
   }
@@ -239,7 +239,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateRoots) {
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
   }
@@ -265,7 +265,7 @@ TEST_F(PGOCtxProfRWTest, DuplicateTargets) {
     ASSERT_TRUE(!!MB);
     ASSERT_NE(*MB, nullptr);
     PGOCtxProfileReader Reader((*MB)->getBuffer());
-    auto Expected = Reader.loadContexts();
+    auto Expected = Reader.loadProfiles();
     EXPECT_FALSE(Expected);
     consumeError(Expected.takeError());
   }