[StandardInstrumentations] Check function analysis invalidation in module passes as well

[aeubanks]

Summary

See comments for why we now need to pass in the MAM instead of the FAM.

[nikic]

LGTM

llvm/include/llvm/Passes/StandardInstrumentations.h

bool invalidate(Function &F, const PreservedAnalyses &PA,
                     FunctionAnalysisManager::Invalidator &);
   };
 
 #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
   SmallVector<StringRef, 8> PassStack;
 #endif
 
   void registerCallbacks(PassInstrumentationCallbacks &PIC,
-                         FunctionAnalysisManager &FAM);
+                         ModuleAnalysisManager &MAM);
 };
 
 // Base class for classes that report changes to the IR.
 // It presents an interface for such classes and provides calls
 // on various events as the new pass manager transforms the IR.
 // It also provides filtering of information based on hidden options
 // specifying which functions are interesting.
 // Calls are made for the following events/queries:

 public:
   StandardInstrumentations(LLVMContext &Context, bool DebugLogging,
                            bool VerifyEach = false,
                            PrintPassOptions PrintPassOpts = PrintPassOptions());
 
   // Register all the standard instrumentation callbacks. If \p FAM is nullptr
   // then PreservedCFGChecker is not enabled.
   void registerCallbacks(PassInstrumentationCallbacks &PIC,
-                         FunctionAnalysisManager *FAM = nullptr);
+                         ModuleAnalysisManager *MAM = nullptr);
 
   TimePassesHandler &getTimePasses() { return TimePasses; }
 };
 
 extern template class ChangeReporter<std::string>;
 extern template class TextChangeReporter<std::string>;
 
 extern template class BlockDataT<EmptyData>;

llvm/lib/LTO/LTOBackend.cpp

static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
 
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
 
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(Mod.getContext(), Conf.DebugPassManager);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
   PassBuilder PB(TM, Conf.PTO, PGOOpt, &PIC);
 
   RegisterPassPlugins(Conf.PassPlugins, PB);
 
   std::unique_ptr<TargetLibraryInfoImpl> TLII(
       new TargetLibraryInfoImpl(Triple(TM->getTargetTriple())));
   if (Conf.Freestanding)
     TLII->disableAllFunctions();

llvm/lib/LTO/ThinLTOCodeGenerator.cpp

static void optimizeModule(Module &TheModule, TargetMachine &TM,
   std::optional<PGOOptions> PGOOpt;
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
 
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(TheModule.getContext(), DebugPassManager);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
   PipelineTuningOptions PTO;
   PTO.LoopVectorization = true;
   PTO.SLPVectorization = true;
   PassBuilder PB(&TM, PTO, PGOOpt, &PIC);
 
   std::unique_ptr<TargetLibraryInfoImpl> TLII(
       new TargetLibraryInfoImpl(Triple(TM.getTargetTriple())));
   if (Freestanding)

llvm/lib/Passes/PassBuilderBindings.cpp

LLVMErrorRef LLVMRunPasses(LLVMModuleRef M, const char *Passes,
   ModuleAnalysisManager MAM;
   PB.registerLoopAnalyses(LAM);
   PB.registerFunctionAnalyses(FAM);
   PB.registerCGSCCAnalyses(CGAM);
   PB.registerModuleAnalyses(MAM);
   PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
 
   StandardInstrumentations SI(Mod->getContext(), Debug, VerifyEach);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
   ModulePassManager MPM;
   if (VerifyEach) {
     MPM.addPass(VerifierPass());
   }
   if (auto Err = PB.parsePassPipeline(MPM, Passes)) {
     return wrap(std::move(Err));
   }
 

llvm/lib/Passes/StandardInstrumentations.cpp

 bool PreservedCFGCheckerInstrumentation::CFG::invalidate(
     Function &F, const PreservedAnalyses &PA,
     FunctionAnalysisManager::Invalidator &) {
   auto PAC = PA.getChecker<PreservedCFGCheckerAnalysis>();
   return !(PAC.preserved() || PAC.preservedSet<AllAnalysesOn<Function>>() ||
            PAC.preservedSet<CFGAnalyses>());
 }
 
+static SmallVector<Function *, 1> GetFunctions(Any IR) {
+  SmallVector<Function *, 1> Functions;
+
+  if (const auto **MaybeF = any_cast<const Function *>(&IR)) {
+    Functions.push_back(*const_cast<Function **>(MaybeF));
+  } else if (const auto **MaybeM = any_cast<const Module *>(&IR)) {
+    for (Function &F : **const_cast<Module **>(MaybeM))
+      Functions.push_back(&F);
+  }
+  return Functions;
+}
+
 void PreservedCFGCheckerInstrumentation::registerCallbacks(
-    PassInstrumentationCallbacks &PIC, FunctionAnalysisManager &FAM) {
+    PassInstrumentationCallbacks &PIC, ModuleAnalysisManager &MAM) {
   if (!VerifyAnalysisInvalidation)
     return;
 
-  FAM.registerPass([&] { return PreservedCFGCheckerAnalysis(); });
-  FAM.registerPass([&] { return PreservedFunctionHashAnalysis(); });
-
-  PIC.registerBeforeNonSkippedPassCallback(
-      [this, &FAM](StringRef P, Any IR) {
+  bool Registered = false;
+  PIC.registerBeforeNonSkippedPassCallback([this, &MAM, Registered](
+                                               StringRef P, Any IR) mutable {
 #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
-        assert(&PassStack.emplace_back(P));
 #endif
-        (void)this;
-        const auto **F = any_cast<const Function *>(&IR);
-        if (!F)
-          return;
 
-        // Make sure a fresh CFG snapshot is available before the pass.
-        FAM.getResult<PreservedCFGCheckerAnalysis>(*const_cast<Function *>(*F));
-        FAM.getResult<PreservedFunctionHashAnalysis>(
-            *const_cast<Function *>(*F));
-      });
+    auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(
+                       *const_cast<Module *>(unwrapModule(IR, /*Force=*/true)))
+                    .getManager();
+    if (!Registered) {
+      FAM.registerPass([&] { return PreservedCFGCheckerAnalysis(); });
+      FAM.registerPass([&] { return PreservedFunctionHashAnalysis(); });
+      Registered = true;
+    }
+
+    for (Function *F : GetFunctions(IR)) {
+      FAM.getResult<PreservedCFGCheckerAnalysis>(*F);
+      FAM.getResult<PreservedFunctionHashAnalysis>(*F);
+    }
 
   PIC.registerAfterPassInvalidatedCallback(
       [this](StringRef P, const PreservedAnalyses &PassPA) {
 #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
         assert(PassStack.pop_back_val() == P &&
                "Before and After callbacks must correspond");
 #endif
         (void)this;
       });
 
-  PIC.registerAfterPassCallback([this, &FAM](StringRef P, Any IR,
+  PIC.registerAfterPassCallback([this, &MAM](StringRef P, Any IR,
                                              const PreservedAnalyses &PassPA) {
 #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
     assert(PassStack.pop_back_val() == P &&
            "Before and After callbacks must correspond");
 #endif
     (void)this;
 
-    const auto **MaybeF = any_cast<const Function *>(&IR);
-    if (!MaybeF)
-      return;
-    Function &F = *const_cast<Function *>(*MaybeF);
+    // We have to get the FAM via the MAM, rather than directly use a passed in
+    // FAM because if MAM has not cached the FAM, it won't invalidate function
+    // analyses in FAM.
+    auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(
+                       *const_cast<Module *>(unwrapModule(IR, /*Force=*/true)))
+                    .getManager();
 
+    for (Function *F : GetFunctions(IR)) {
       if (auto *HashBefore =
-            FAM.getCachedResult<PreservedFunctionHashAnalysis>(F)) {
-      if (HashBefore->Hash != StructuralHash(F)) {
+              FAM.getCachedResult<PreservedFunctionHashAnalysis>(*F)) {
+        if (HashBefore->Hash != StructuralHash(*F)) {
           report_fatal_error(formatv(
               "Function @{0} changed by {1} without invalidating analyses",
-            F.getName(), P));
+              F->getName(), P));
         }
       }
 
       auto CheckCFG = [](StringRef Pass, StringRef FuncName,
                          const CFG &GraphBefore, const CFG &GraphAfter) {
         if (GraphAfter == GraphBefore)
           return;
 
-      dbgs() << "Error: " << Pass
-             << " does not invalidate CFG analyses but CFG changes detected in "
-                "function @"
-             << FuncName << ":\n";
+        dbgs()
+            << "Error: " << Pass
         CFG::printDiff(dbgs(), GraphBefore, GraphAfter);
         report_fatal_error(Twine("CFG unexpectedly changed by ", Pass));
       };
 
-    if (auto *GraphBefore = FAM.getCachedResult<PreservedCFGCheckerAnalysis>(F))
-      CheckCFG(P, F.getName(), *GraphBefore,
-               CFG(&F, /* TrackBBLifetime */ false));
+      if (auto *GraphBefore =
+              FAM.getCachedResult<PreservedCFGCheckerAnalysis>(*F))
+        CheckCFG(P, F->getName(), *GraphBefore,
+                 CFG(F, /* TrackBBLifetime */ false));
+    }
   });
 }
 
 void VerifyInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   PIC.registerAfterPassCallback(
       [this](StringRef P, Any IR, const PreservedAnalyses &PassPA) {
         if (isIgnored(P) || P == "VerifierPass")
if (!isInteresting(IR, PassID, PIC.getPassNameForClassName(PassID))) {
       return;
     }
     OS << " Started ***\n";
     unwrapAndPrint(OS, IR);
   });
 }
 
 void StandardInstrumentations::registerCallbacks(
-    PassInstrumentationCallbacks &PIC, FunctionAnalysisManager *FAM) {
+    PassInstrumentationCallbacks &PIC, ModuleAnalysisManager *MAM) {
   PrintIR.registerCallbacks(PIC);
   PrintPass.registerCallbacks(PIC);
   TimePasses.registerCallbacks(PIC);
   OptNone.registerCallbacks(PIC);
   OptPassGate.registerCallbacks(PIC);
   PrintChangedIR.registerCallbacks(PIC);
   PseudoProbeVerification.registerCallbacks(PIC);
   if (VerifyEach)
     Verify.registerCallbacks(PIC);
   PrintChangedDiff.registerCallbacks(PIC);
   WebsiteChangeReporter.registerCallbacks(PIC);
   ChangeTester.registerCallbacks(PIC);
   PrintCrashIR.registerCallbacks(PIC);
-  if (FAM)
-    PreservedCFGChecker.registerCallbacks(PIC, *FAM);
+  if (MAM)
+    PreservedCFGChecker.registerCallbacks(PIC, *MAM);
 
   // TimeProfiling records the pass running time cost.
   // Its 'BeforePassCallback' can be appended at the tail of all the
   // BeforeCallbacks by calling `registerCallbacks` in the end.
   // Its 'AfterPassCallback' is put at the front of all the
   // AfterCallbacks by its `registerCallbacks`. This is necessary
   // to ensure that other callbacks are not included in the timings.
   TimeProfilingPasses.registerCallbacks(PIC);

llvm/tools/opt/NewPMDriver.cpp

bool llvm::runPassPipeline(StringRef Arg0, Module &M, TargetMachine *TM,
   ModuleAnalysisManager MAM;
 
   PassInstrumentationCallbacks PIC;
   PrintPassOptions PrintPassOpts;
   PrintPassOpts.Verbose = DebugPM == DebugLogging::Verbose;
   PrintPassOpts.SkipAnalyses = DebugPM == DebugLogging::Quiet;
   StandardInstrumentations SI(M.getContext(), DebugPM != DebugLogging::None,
                               VerifyEachPass, PrintPassOpts);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
   DebugifyEachInstrumentation Debugify;
   DebugifyStatsMap DIStatsMap;
   DebugInfoPerPass DebugInfoBeforePass;
   if (DebugifyEach) {
     Debugify.setDIStatsMap(DIStatsMap);
     Debugify.setDebugifyMode(DebugifyMode::SyntheticDebugInfo);
     Debugify.registerCallbacks(PIC, FAM);
   } else if (VerifyEachDebugInfoPreserve) {

llvm/unittests/IR/PassManagerTest.cpp

auto M = parseIR(Context, "define void @foo() {\n"
                             "next:\n"
                             "  br label %exit\n"
                             "exit:\n"
                             "  ret void\n"
                             "}\n");
 
   auto *F = M->getFunction("foo");
   FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
   FunctionPassManager FPM;
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ true);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return DominatorTreeAnalysis(); });
   FAM.registerPass([&] { return AssumptionAnalysis(); });
   FAM.registerPass([&] { return TargetIRAnalysis(); });
 
   FPM.addPass(SimplifyCFGPass());
   FPM.run(*F, FAM);
 }
auto M = parseIR(Context, "define void @foo() {\n"
                             "next:\n"
                             "  br label %exit\n"
                             "exit:\n"
                             "  ret void\n"
                             "}\n");
 
   auto *F = M->getFunction("foo");
   FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
   FunctionPassManager FPM;
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ true);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return DominatorTreeAnalysis(); });
   FAM.registerPass([&] { return AssumptionAnalysis(); });
   FAM.registerPass([&] { return TargetIRAnalysis(); });
 
   FPM.addPass(TestSimplifyCFGInvalidatingAnalysisPass());
   FPM.run(*F, FAM);
 }
auto M = parseIR(Context, "define void @foo() {\n"
                             "next:\n"
                             "  br label %exit\n"
                             "exit:\n"
                             "  ret void\n"
                             "}\n");
 
   auto *F = M->getFunction("foo");
   FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
   FunctionPassManager FPM;
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ true);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return DominatorTreeAnalysis(); });
   FAM.registerPass([&] { return AssumptionAnalysis(); });
   FAM.registerPass([&] { return TargetIRAnalysis(); });
 
   FunctionPassManager InnerFPM;
   InnerFPM.addPass(SimplifyCFGPass());
 
   FPM.addPass(TestSimplifyCFGWrapperPass(InnerFPM));
   FPM.run(*F, FAM);
 }
 
 #ifdef EXPENSIVE_CHECKS
 
 struct WrongFunctionPass : PassInfoMixin<WrongFunctionPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM) {
     F.getEntryBlock().begin()->eraseFromParent();
     return PreservedAnalyses::all();
   }
   static StringRef name() { return "WrongFunctionPass"; }
 };
 
-TEST_F(PassManagerTest, FunctionAnalysisMissedInvalidation) {
+TEST_F(PassManagerTest, FunctionPassMissedFunctionAnalysisInvalidation) {
   LLVMContext Context;
   auto M = parseIR(Context, "define void @foo() {\n"
                             "  %a = add i32 0, 0\n"
                             "  ret void\n"
                             "}\n");
 
   FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
   PassInstrumentationCallbacks PIC;
   StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ false);
-  SI.registerCallbacks(PIC, &FAM);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
   FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
 
   FunctionPassManager FPM;
   FPM.addPass(WrongFunctionPass());
 
   auto *F = M->getFunction("foo");
   EXPECT_DEATH(FPM.run(*F, FAM), "Function @foo changed by WrongFunctionPass without invalidating analyses");
 }
 
-#endif
+struct WrongModulePass : PassInfoMixin<WrongModulePass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    for (Function &F : M)
+      F.getEntryBlock().begin()->eraseFromParent();
 
+    return PreservedAnalyses::all();
+  }
+  static StringRef name() { return "WrongModulePass"; }
+};
+
+TEST_F(PassManagerTest, ModulePassMissedFunctionAnalysisInvalidation) {
+  LLVMContext Context;
+  auto M = parseIR(Context, "define void @foo() {\n"
+                            "  %a = add i32 0, 0\n"
+                            "  ret void\n"
+                            "}\n");
+
+  FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
+  PassInstrumentationCallbacks PIC;
+  StandardInstrumentations SI(M->getContext(), /*DebugLogging*/ false);
+  SI.registerCallbacks(PIC, &MAM);
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+  FAM.registerPass([&] { return PassInstrumentationAnalysis(&PIC); });
+
+  ModulePassManager MPM;
+  MPM.addPass(WrongModulePass());
+
+  EXPECT_DEATH(
+      MPM.run(*M, MAM),
+      "Function @foo changed by WrongModulePass without invalidating analyses");
+}
+
+#endif
 }