[NewPM] Set diagnostic handler in MachineModuleAnalysis (#88229)

`setDiagnosticHandler` is idempotent.

Author: paperchalice

llvm/lib/CodeGen/MachineModuleInfo.cpp

@@ -213,7 +213,6 @@ static unsigned getLocCookie(const SMDiagnostic &SMD, const SourceMgr &SrcMgr,
 bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {
   MMI.initialize();
   MMI.TheModule = &M;
-  // FIXME: Do this for new pass manager.
   LLVMContext &Ctx = M.getContext();
   MMI.getContext().setDiagnosticHandler(
       [&Ctx, &M](const SMDiagnostic &SMD, bool IsInlineAsm,
@@ -240,6 +239,17 @@ AnalysisKey MachineModuleAnalysis::Key;
 MachineModuleAnalysis::Result
 MachineModuleAnalysis::run(Module &M, ModuleAnalysisManager &) {
   MMI.TheModule = &M;
+  LLVMContext &Ctx = M.getContext();
+  MMI.getContext().setDiagnosticHandler(
+      [&Ctx, &M](const SMDiagnostic &SMD, bool IsInlineAsm,
+                 const SourceMgr &SrcMgr,
+                 std::vector<const MDNode *> &LocInfos) {
+        unsigned LocCookie = 0;
+        if (IsInlineAsm)
+          LocCookie = getLocCookie(SMD, SrcMgr, LocInfos);
+        Ctx.diagnose(
+            DiagnosticInfoSrcMgr(SMD, M.getName(), IsInlineAsm, LocCookie));
+      });
   MMI.DbgInfoAvailable =
       !DisableDebugInfoPrinting && !M.debug_compile_units().empty();
   return Result(MMI);

llvm/unittests/CodeGen/PassManagerTest.cpp

@@ -125,6 +125,15 @@ struct TestMachineModulePass : public PassInfoMixin<TestMachineModulePass> {
   std::vector<int> &Counts;
 };
 
+struct ReportWarningPass : public PassInfoMixin<ReportWarningPass> {
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM) {
+    auto &Ctx = MF.getContext();
+    Ctx.reportWarning(SMLoc(), "Test warning message.");
+    return PreservedAnalyses::all();
+  }
+};
+
 std::unique_ptr<Module> parseIR(LLVMContext &Context, const char *IR) {
   SMDiagnostic Err;
   return parseAssemblyString(IR, Err, Context);
@@ -176,10 +185,10 @@ TEST_F(PassManagerTest, Basic) {
   MachineModuleInfo MMI(LLVMTM);
 
   LoopAnalysisManager LAM;
+  MachineFunctionAnalysisManager MFAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
-  MachineFunctionAnalysisManager MFAM;
   PassBuilder PB(TM.get());
   PB.registerModuleAnalyses(MAM);
   PB.registerCGSCCAnalyses(CGAM);
@@ -204,10 +213,49 @@ TEST_F(PassManagerTest, Basic) {
   MFPM.addPass(TestMachineFunctionPass(Count, Counts));
   MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
 
+  testing::internal::CaptureStderr();
   MPM.run(*M, MAM);
+  std::string Output = testing::internal::GetCapturedStderr();
 
   EXPECT_EQ((std::vector<int>{10, 16, 18, 20, 30, 36, 38, 40}), Counts);
   EXPECT_EQ(40, Count);
 }
 
+TEST_F(PassManagerTest, DiagnosticHandler) {
+  if (!TM)
+    GTEST_SKIP();
+
+  LLVMTargetMachine *LLVMTM = static_cast<LLVMTargetMachine *>(TM.get());
+  M->setDataLayout(TM->createDataLayout());
+
+  MachineModuleInfo MMI(LLVMTM);
+
+  LoopAnalysisManager LAM;
+  MachineFunctionAnalysisManager MFAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  PassBuilder PB(TM.get());
+  PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.registerMachineFunctionAnalyses(MFAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM, &MFAM);
+
+  MAM.registerPass([&] { return MachineModuleAnalysis(MMI); });
+
+  ModulePassManager MPM;
+  FunctionPassManager FPM;
+  MachineFunctionPassManager MFPM;
+  MFPM.addPass(ReportWarningPass());
+  MPM.addPass(createModuleToMachineFunctionPassAdaptor(std::move(MFPM)));
+  testing::internal::CaptureStderr();
+  MPM.run(*M, MAM);
+  std::string Output = testing::internal::GetCapturedStderr();
+
+  EXPECT_TRUE(Output.find("warning: <unknown>:0: Test warning message.") !=
+              std::string::npos);
+}
+
 } // namespace