[StandardInstrumentations] Check that module analyses are properly invalidated

[aeubanks]

Summary

Followup to D146003/D146160

[nikic]

llvm/lib/Passes/StandardInstrumentations.cpp

1126

Do this outside the pass callback?

1204

Hm, does that mean that in most cases we're going to verify both all functions via the function hash, and then again using the module hash (which includes all function hashes)?

[aeubanks]

llvm/lib/Passes/StandardInstrumentations.cpp

1126

might as well be consistent with the FAM code, which requires having some IR to get from the MAM (or we coul pass in the FAM, but that's more work for callers)

1204

yes, but that's fine, it's only doubling this work (which shouldn't matter in expensive checks builds)

we could have the module analysis fetch the hash of each function via the function hash analysis, but at that point we start mixing module and function analyses which I'd like to avoid to keep this conceptually simple and not fall into traps with cross-level analyses (and it would require changes of how we hash modules by combining function hashes rather than hashing everything one after the other, but that's not too much work)

[nikic]

LG

llvm/lib/Passes/StandardInstrumentations.cpp

1204

Okay, let's keep it simple for now...

llvm/lib/Passes/StandardInstrumentations.cpp

struct PreservedFunctionHashAnalysis
 
   Result run(Function &F, FunctionAnalysisManager &FAM) {
     return Result{StructuralHash(F)};
   }
 };
 
 AnalysisKey PreservedFunctionHashAnalysis::Key;
 
+struct PreservedModuleHashAnalysis
+    : public AnalysisInfoMixin<PreservedModuleHashAnalysis> {
+  static AnalysisKey Key;
+
+  struct ModuleHash {
+    uint64_t Hash;
+  };
+
+  using Result = ModuleHash;
+
+  Result run(Module &F, ModuleAnalysisManager &FAM) {
+    return Result{StructuralHash(F)};
+  }
+};
+
+AnalysisKey PreservedModuleHashAnalysis::Key;
+
 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>());
 }
 
#endif
     (void)this;
 
     auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(
                        *const_cast<Module *>(unwrapModule(IR, /*Force=*/true)))
                     .getManager();
     if (!Registered) {
       FAM.registerPass([&] { return PreservedCFGCheckerAnalysis(); });
       FAM.registerPass([&] { return PreservedFunctionHashAnalysis(); });
+      MAM.registerPass([&] { return PreservedModuleHashAnalysis(); });
       Registered = true;
     }
 
     for (Function *F : GetFunctions(IR)) {
       // Make sure a fresh CFG snapshot is available before the pass.
       FAM.getResult<PreservedCFGCheckerAnalysis>(*F);
       FAM.getResult<PreservedFunctionHashAnalysis>(*F);
     }
+
+    if (auto *MaybeM = any_cast<const Module *>(&IR)) {
+      Module &M = **const_cast<Module **>(MaybeM);
+      MAM.getResult<PreservedModuleHashAnalysis>(M);
+    }
   });
 
   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
for (Function *F : GetFunctions(IR)) {
         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 *MaybeM = any_cast<const Module *>(&IR)) {
+      Module &M = **const_cast<Module **>(MaybeM);
+      if (auto *HashBefore =
+              MAM.getCachedResult<PreservedModuleHashAnalysis>(M)) {
+        if (HashBefore->Hash != StructuralHash(M)) {
+          report_fatal_error(formatv(
+              "Module changed by {0} without invalidating analyses", P));
+        }
+      }
+    }
   });
 }
 
 void VerifyInstrumentation::registerCallbacks(
     PassInstrumentationCallbacks &PIC) {
   PIC.registerAfterPassCallback(
       [this](StringRef P, Any IR, const PreservedAnalyses &PassPA) {
         if (isIgnored(P) || P == "VerifierPass")

llvm/unittests/IR/PassManagerTest.cpp

TEST_F(PassManagerTest, FunctionPassMissedFunctionAnalysisInvalidation) {
   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); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
 
   FunctionPassManager FPM;
   FPM.addPass(WrongFunctionPass());
 
   auto *F = M->getFunction("foo");
   EXPECT_DEATH(FPM.run(*F, FAM), "Function @foo changed by WrongFunctionPass without invalidating analyses");
 }
 
 struct WrongModulePass : PassInfoMixin<WrongModulePass> {
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
     for (Function &F : M)
       F.getEntryBlock().begin()->eraseFromParent();
 
-    return PreservedAnalyses::all();
+    PreservedAnalyses PA;
+    PA.preserveSet<AllAnalysesOn<Function>>();
+    PA.preserve<FunctionAnalysisManagerModuleProxy>();
+    return PA;
   }
   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); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
 
   ModulePassManager MPM;
   MPM.addPass(WrongModulePass());
 
   EXPECT_DEATH(
       MPM.run(*M, MAM),
       "Function @foo changed by WrongModulePass without invalidating analyses");
 }
 
+struct WrongModulePass2 : PassInfoMixin<WrongModulePass2> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
+    for (Function &F : M)
+      F.getEntryBlock().begin()->eraseFromParent();
+
+    PreservedAnalyses PA;
+    PA.preserveSet<AllAnalysesOn<Module>>();
+    PA.abandon<FunctionAnalysisManagerModuleProxy>();
+    return PA;
+  }
+  static StringRef name() { return "WrongModulePass2"; }
+};
+
+TEST_F(PassManagerTest, ModulePassMissedModuleAnalysisInvalidation) {
+  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); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
+
+  ModulePassManager MPM;
+  MPM.addPass(WrongModulePass2());
+
+  EXPECT_DEATH(
+      MPM.run(*M, MAM),
+      "Module changed by WrongModulePass2 without invalidating analyses");
+}
+
 #endif
 }