1 files changed, 6 insertions, 17 deletions
diff --git a/llvm/lib/Transforms/Instrumentation/CGProfile.cpp b/llvm/lib/Transforms/Instrumentation/CGProfile.cpp
index 1c630e9ee424..d53e12ad1ff5 100644
--- a/llvm/lib/Transforms/Instrumentation/CGProfile.cpp
+++ b/llvm/lib/Transforms/Instrumentation/CGProfile.cpp
@@ -15,7 +15,6 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/PassManager.h"
-#include "llvm/InitializePasses.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Transforms/Instrumentation.h"
 #include <optional>
@@ -46,8 +45,7 @@ addModuleFlags(Module &M,
 }
 
 static bool runCGProfilePass(
-    Module &M, function_ref<BlockFrequencyInfo &(Function &)> GetBFI,
-    function_ref<TargetTransformInfo &(Function &)> GetTTI, bool LazyBFI) {
+    Module &M, FunctionAnalysisManager &FAM) {
   MapVector<std::pair<Function *, Function *>, uint64_t> Counts;
   InstrProfSymtab Symtab;
   auto UpdateCounts = [&](TargetTransformInfo &TTI, Function *F,
@@ -64,15 +62,13 @@ static bool runCGProfilePass(
   (void)(bool) Symtab.create(M);
   for (auto &F : M) {
     // Avoid extra cost of running passes for BFI when the function doesn't have
-    // entry count. Since LazyBlockFrequencyInfoPass only exists in LPM, check
-    // if using LazyBlockFrequencyInfoPass.
-    // TODO: Remove LazyBFI when LazyBlockFrequencyInfoPass is available in NPM.
-    if (F.isDeclaration() || (LazyBFI && !F.getEntryCount()))
+    // entry count.
+    if (F.isDeclaration() || !F.getEntryCount())
       continue;
-    auto &BFI = GetBFI(F);
+    auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
     if (BFI.getEntryFreq() == 0)
       continue;
-    TargetTransformInfo &TTI = GetTTI(F);
+    TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
     for (auto &BB : F) {
       std::optional<uint64_t> BBCount = BFI.getBlockProfileCount(&BB);
       if (!BBCount)
@@ -105,14 +101,7 @@ static bool runCGProfilePass(
 PreservedAnalyses CGProfilePass::run(Module &M, ModuleAnalysisManager &MAM) {
   FunctionAnalysisManager &FAM =
       MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-  auto GetBFI = [&FAM](Function &F) -> BlockFrequencyInfo & {
-    return FAM.getResult<BlockFrequencyAnalysis>(F);
-  };
-  auto GetTTI = [&FAM](Function &F) -> TargetTransformInfo & {
-    return FAM.getResult<TargetIRAnalysis>(F);
-  };
-
-  runCGProfilePass(M, GetBFI, GetTTI, false);
+  runCGProfilePass(M, FAM);
 
   return PreservedAnalyses::all();
 }