[CodeGen] Implement post-opt linking option for builtin bitocdes (#69371)

In this patch, we create a new ModulePass that mimics the LinkInModules
API from CodeGenAction.cpp, and a new command line option to enable the
pass. As part of the implementation, we needed to refactor the
BackendConsumer class definition into a new separate header (instead of
embedded in CodeGenAction.cpp). With this new pass, we can now re-link
bitcodes supplied via the -mlink-built-in bitcodes as part of the
RunOptimizationPipeline.

With the re-linking pass, we now handle cases where new device library
functions are introduced as part of the optimization pipeline.
Previously, these newly introduced functions (for example a fused sincos
call) would result in a linking error due to a missing function
definition. This new pass can be initiated via:

      -mllvm -relink-builtin-bitcode-postop

Also note we intentionally exclude bitcodes supplied via the
-mlink-bitcode-file option from the second linking step

1 files changed, 299 insertions, 361 deletions

diff --git a/clang/lib/CodeGen/CodeGenAction.cpp b/clang/lib/CodeGen/CodeGenAction.cpp
index a3b7238..a31a271 100644
--- a/clang/lib/CodeGen/CodeGenAction.cpp
+++ b/clang/lib/CodeGen/CodeGenAction.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/CodeGen/CodeGenAction.h"
+#include "BackendConsumer.h"
 #include "CGCall.h"
 #include "CodeGenModule.h"
 #include "CoverageMappingGen.h"
@@ -48,8 +49,8 @@
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Transforms/IPO/Internalize.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 
-#include <memory>
 #include <optional>
 using namespace clang;
 using namespace llvm;
@@ -57,419 +58,356 @@ using namespace llvm;
 #define DEBUG_TYPE "codegenaction"
 
 namespace clang {
-  class BackendConsumer;
-  class ClangDiagnosticHandler final : public DiagnosticHandler {
-  public:
-    ClangDiagnosticHandler(const CodeGenOptions &CGOpts, BackendConsumer *BCon)
-        : CodeGenOpts(CGOpts), BackendCon(BCon) {}
+class BackendConsumer;
+class ClangDiagnosticHandler final : public DiagnosticHandler {
+public:
+  ClangDiagnosticHandler(const CodeGenOptions &CGOpts, BackendConsumer *BCon)
+      : CodeGenOpts(CGOpts), BackendCon(BCon) {}
 
-    bool handleDiagnostics(const DiagnosticInfo &DI) override;
+  bool handleDiagnostics(const DiagnosticInfo &DI) override;
 
-    bool isAnalysisRemarkEnabled(StringRef PassName) const override {
-      return CodeGenOpts.OptimizationRemarkAnalysis.patternMatches(PassName);
-    }
-    bool isMissedOptRemarkEnabled(StringRef PassName) const override {
-      return CodeGenOpts.OptimizationRemarkMissed.patternMatches(PassName);
-    }
-    bool isPassedOptRemarkEnabled(StringRef PassName) const override {
-      return CodeGenOpts.OptimizationRemark.patternMatches(PassName);
-    }
+  bool isAnalysisRemarkEnabled(StringRef PassName) const override {
+    return CodeGenOpts.OptimizationRemarkAnalysis.patternMatches(PassName);
+  }
+  bool isMissedOptRemarkEnabled(StringRef PassName) const override {
+    return CodeGenOpts.OptimizationRemarkMissed.patternMatches(PassName);
+  }
+  bool isPassedOptRemarkEnabled(StringRef PassName) const override {
+    return CodeGenOpts.OptimizationRemark.patternMatches(PassName);
+  }
 
-    bool isAnyRemarkEnabled() const override {
-      return CodeGenOpts.OptimizationRemarkAnalysis.hasValidPattern() ||
-             CodeGenOpts.OptimizationRemarkMissed.hasValidPattern() ||
-             CodeGenOpts.OptimizationRemark.hasValidPattern();
-    }
+  bool isAnyRemarkEnabled() const override {
+    return CodeGenOpts.OptimizationRemarkAnalysis.hasValidPattern() ||
+           CodeGenOpts.OptimizationRemarkMissed.hasValidPattern() ||
+           CodeGenOpts.OptimizationRemark.hasValidPattern();
+  }
 
-  private:
-    const CodeGenOptions &CodeGenOpts;
-    BackendConsumer *BackendCon;
-  };
+private:
+  const CodeGenOptions &CodeGenOpts;
+  BackendConsumer *BackendCon;
+};
+
+static void reportOptRecordError(Error E, DiagnosticsEngine &Diags,
+                                 const CodeGenOptions &CodeGenOpts) {
+  handleAllErrors(
+      std::move(E),
+    [&](const LLVMRemarkSetupFileError &E) {
+        Diags.Report(diag::err_cannot_open_file)
+            << CodeGenOpts.OptRecordFile << E.message();
+      },
+    [&](const LLVMRemarkSetupPatternError &E) {
+        Diags.Report(diag::err_drv_optimization_remark_pattern)
+            << E.message() << CodeGenOpts.OptRecordPasses;
+      },
+    [&](const LLVMRemarkSetupFormatError &E) {
+        Diags.Report(diag::err_drv_optimization_remark_format)
+            << CodeGenOpts.OptRecordFormat;
+      });
+}
 
-  static void reportOptRecordError(Error E, DiagnosticsEngine &Diags,
-                                   const CodeGenOptions &CodeGenOpts) {
-    handleAllErrors(
-        std::move(E),
-      [&](const LLVMRemarkSetupFileError &E) {
-          Diags.Report(diag::err_cannot_open_file)
-              << CodeGenOpts.OptRecordFile << E.message();
-        },
-      [&](const LLVMRemarkSetupPatternError &E) {
-          Diags.Report(diag::err_drv_optimization_remark_pattern)
-              << E.message() << CodeGenOpts.OptRecordPasses;
-        },
-      [&](const LLVMRemarkSetupFormatError &E) {
-          Diags.Report(diag::err_drv_optimization_remark_format)
-              << CodeGenOpts.OptRecordFormat;
-        });
-    }
+BackendConsumer::BackendConsumer(BackendAction Action, DiagnosticsEngine &Diags,
+                                 IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+                                 const HeaderSearchOptions &HeaderSearchOpts,
+                                 const PreprocessorOptions &PPOpts,
+                                 const CodeGenOptions &CodeGenOpts,
+                                 const TargetOptions &TargetOpts,
+                                 const LangOptions &LangOpts,
+                                 const std::string &InFile,
+                                 SmallVector<LinkModule, 4> LinkModules,
+                                 std::unique_ptr<raw_pwrite_stream> OS,
+                                 LLVMContext &C,
+                                 CoverageSourceInfo *CoverageInfo)
+  : Diags(Diags), Action(Action), HeaderSearchOpts(HeaderSearchOpts),
+  CodeGenOpts(CodeGenOpts), TargetOpts(TargetOpts), LangOpts(LangOpts),
+  AsmOutStream(std::move(OS)), Context(nullptr), FS(VFS),
+  LLVMIRGeneration("irgen", "LLVM IR Generation Time"),
+  LLVMIRGenerationRefCount(0),
+  Gen(CreateLLVMCodeGen(Diags, InFile, std::move(VFS), HeaderSearchOpts,
+                        PPOpts, CodeGenOpts, C, CoverageInfo)),
+  LinkModules(std::move(LinkModules)) {
+    TimerIsEnabled = CodeGenOpts.TimePasses;
+    llvm::TimePassesIsEnabled = CodeGenOpts.TimePasses;
+    llvm::TimePassesPerRun = CodeGenOpts.TimePassesPerRun;
+}
 
-  class BackendConsumer : public ASTConsumer {
-    using LinkModule = CodeGenAction::LinkModule;
-
-    virtual void anchor();
-    DiagnosticsEngine &Diags;
-    BackendAction Action;
-    const HeaderSearchOptions &HeaderSearchOpts;
-    const CodeGenOptions &CodeGenOpts;
-    const TargetOptions &TargetOpts;
-    const LangOptions &LangOpts;
-    std::unique_ptr<raw_pwrite_stream> AsmOutStream;
-    ASTContext *Context;
-    IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS;
-
-    Timer LLVMIRGeneration;
-    unsigned LLVMIRGenerationRefCount;
-
-    /// True if we've finished generating IR. This prevents us from generating
-    /// additional LLVM IR after emitting output in HandleTranslationUnit. This
-    /// can happen when Clang plugins trigger additional AST deserialization.
-    bool IRGenFinished = false;
-
-    bool TimerIsEnabled = false;
-
-    std::unique_ptr<CodeGenerator> Gen;
-
-    SmallVector<LinkModule, 4> LinkModules;
-
-    // A map from mangled names to their function's source location, used for
-    // backend diagnostics as the Clang AST may be unavailable. We actually use
-    // the mangled name's hash as the key because mangled names can be very
-    // long and take up lots of space. Using a hash can cause name collision,
-    // but that is rare and the consequences are pointing to a wrong source
-    // location which is not severe. This is a vector instead of an actual map
-    // because we optimize for time building this map rather than time
-    // retrieving an entry, as backend diagnostics are uncommon.
-    std::vector<std::pair<llvm::hash_code, FullSourceLoc>>
-        ManglingFullSourceLocs;
-
-    // This is here so that the diagnostic printer knows the module a diagnostic
-    // refers to.
-    llvm::Module *CurLinkModule = nullptr;
-
-  public:
-    BackendConsumer(BackendAction Action, DiagnosticsEngine &Diags,
-                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
-                    const HeaderSearchOptions &HeaderSearchOpts,
-                    const PreprocessorOptions &PPOpts,
-                    const CodeGenOptions &CodeGenOpts,
-                    const TargetOptions &TargetOpts,
-                    const LangOptions &LangOpts, const std::string &InFile,
-                    SmallVector<LinkModule, 4> LinkModules,
-                    std::unique_ptr<raw_pwrite_stream> OS, LLVMContext &C,
-                    CoverageSourceInfo *CoverageInfo = nullptr)
-        : Diags(Diags), Action(Action), HeaderSearchOpts(HeaderSearchOpts),
-          CodeGenOpts(CodeGenOpts), TargetOpts(TargetOpts), LangOpts(LangOpts),
-          AsmOutStream(std::move(OS)), Context(nullptr), FS(VFS),
-          LLVMIRGeneration("irgen", "LLVM IR Generation Time"),
-          LLVMIRGenerationRefCount(0),
-          Gen(CreateLLVMCodeGen(Diags, InFile, std::move(VFS), HeaderSearchOpts,
-                                PPOpts, CodeGenOpts, C, CoverageInfo)),
-          LinkModules(std::move(LinkModules)) {
-      TimerIsEnabled = CodeGenOpts.TimePasses;
-      llvm::TimePassesIsEnabled = CodeGenOpts.TimePasses;
-      llvm::TimePassesPerRun = CodeGenOpts.TimePassesPerRun;
-    }
+// This constructor is used in installing an empty BackendConsumer
+// to use the clang diagnostic handler for IR input files. It avoids
+// initializing the OS field.
+BackendConsumer::BackendConsumer(BackendAction Action, DiagnosticsEngine &Diags,
+                                 IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+                                 const HeaderSearchOptions &HeaderSearchOpts,
+                                 const PreprocessorOptions &PPOpts,
+                                 const CodeGenOptions &CodeGenOpts,
+                                 const TargetOptions &TargetOpts,
+                                 const LangOptions &LangOpts,
+                                 llvm::Module *Module,
+                                 SmallVector<LinkModule, 4> LinkModules,
+                                 LLVMContext &C,
+                                 CoverageSourceInfo *CoverageInfo)
+  : Diags(Diags), Action(Action), HeaderSearchOpts(HeaderSearchOpts),
+  CodeGenOpts(CodeGenOpts), TargetOpts(TargetOpts), LangOpts(LangOpts),
+  Context(nullptr), FS(VFS),
+  LLVMIRGeneration("irgen", "LLVM IR Generation Time"),
+  LLVMIRGenerationRefCount(0),
+  Gen(CreateLLVMCodeGen(Diags, "", std::move(VFS), HeaderSearchOpts,
+                        PPOpts, CodeGenOpts, C, CoverageInfo)),
+  LinkModules(std::move(LinkModules)), CurLinkModule(Module) {
+    TimerIsEnabled = CodeGenOpts.TimePasses;
+    llvm::TimePassesIsEnabled = CodeGenOpts.TimePasses;
+    llvm::TimePassesPerRun = CodeGenOpts.TimePassesPerRun;
+}
 
-    // This constructor is used in installing an empty BackendConsumer
-    // to use the clang diagnostic handler for IR input files. It avoids
-    // initializing the OS field.
-    BackendConsumer(BackendAction Action, DiagnosticsEngine &Diags,
-                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
-                    const HeaderSearchOptions &HeaderSearchOpts,
-                    const PreprocessorOptions &PPOpts,
-                    const CodeGenOptions &CodeGenOpts,
-                    const TargetOptions &TargetOpts,
-                    const LangOptions &LangOpts, llvm::Module *Module,
-                    SmallVector<LinkModule, 4> LinkModules, LLVMContext &C,
-                    CoverageSourceInfo *CoverageInfo = nullptr)
-        : Diags(Diags), Action(Action), HeaderSearchOpts(HeaderSearchOpts),
-          CodeGenOpts(CodeGenOpts), TargetOpts(TargetOpts), LangOpts(LangOpts),
-          Context(nullptr), FS(VFS),
-          LLVMIRGeneration("irgen", "LLVM IR Generation Time"),
-          LLVMIRGenerationRefCount(0),
-          Gen(CreateLLVMCodeGen(Diags, "", std::move(VFS), HeaderSearchOpts,
-                                PPOpts, CodeGenOpts, C, CoverageInfo)),
-          LinkModules(std::move(LinkModules)), CurLinkModule(Module) {
-      TimerIsEnabled = CodeGenOpts.TimePasses;
-      llvm::TimePassesIsEnabled = CodeGenOpts.TimePasses;
-      llvm::TimePassesPerRun = CodeGenOpts.TimePassesPerRun;
-    }
-    llvm::Module *getModule() const { return Gen->GetModule(); }
-    std::unique_ptr<llvm::Module> takeModule() {
-      return std::unique_ptr<llvm::Module>(Gen->ReleaseModule());
-    }
+llvm::Module* BackendConsumer::getModule() const {
+  return Gen->GetModule();
+}
 
-    CodeGenerator *getCodeGenerator() { return Gen.get(); }
+std::unique_ptr<llvm::Module> BackendConsumer::takeModule() {
+  return std::unique_ptr<llvm::Module>(Gen->ReleaseModule());
+}
 
-    void HandleCXXStaticMemberVarInstantiation(VarDecl *VD) override {
-      Gen->HandleCXXStaticMemberVarInstantiation(VD);
-    }
+CodeGenerator* BackendConsumer::getCodeGenerator() {
+  return Gen.get();
+}
 
-    void Initialize(ASTContext &Ctx) override {
-      assert(!Context && "initialized multiple times");
+void BackendConsumer::HandleCXXStaticMemberVarInstantiation(VarDecl *VD) {
+  Gen->HandleCXXStaticMemberVarInstantiation(VD);
+}
 
-      Context = &Ctx;
+void BackendConsumer::Initialize(ASTContext &Ctx) {
+  assert(!Context && "initialized multiple times");
 
-      if (TimerIsEnabled)
-        LLVMIRGeneration.startTimer();
+  Context = &Ctx;
 
-      Gen->Initialize(Ctx);
+  if (TimerIsEnabled)
+    LLVMIRGeneration.startTimer();
 
-      if (TimerIsEnabled)
-        LLVMIRGeneration.stopTimer();
-    }
+  Gen->Initialize(Ctx);
 
-    bool HandleTopLevelDecl(DeclGroupRef D) override {
-      PrettyStackTraceDecl CrashInfo(*D.begin(), SourceLocation(),
-                                     Context->getSourceManager(),
-                                     "LLVM IR generation of declaration");
+  if (TimerIsEnabled)
+    LLVMIRGeneration.stopTimer();
+}
 
-      // Recurse.
-      if (TimerIsEnabled) {
-        LLVMIRGenerationRefCount += 1;
-        if (LLVMIRGenerationRefCount == 1)
-          LLVMIRGeneration.startTimer();
-      }
+bool BackendConsumer::HandleTopLevelDecl(DeclGroupRef D) {
+  PrettyStackTraceDecl CrashInfo(*D.begin(), SourceLocation(),
+                                 Context->getSourceManager(),
+                                 "LLVM IR generation of declaration");
 
-      Gen->HandleTopLevelDecl(D);
+  // Recurse.
+  if (TimerIsEnabled) {
+    LLVMIRGenerationRefCount += 1;
+    if (LLVMIRGenerationRefCount == 1)
+      LLVMIRGeneration.startTimer();
+  }
 
-      if (TimerIsEnabled) {
-        LLVMIRGenerationRefCount -= 1;
-        if (LLVMIRGenerationRefCount == 0)
-          LLVMIRGeneration.stopTimer();
-      }
+  Gen->HandleTopLevelDecl(D);
 
-      return true;
-    }
+  if (TimerIsEnabled) {
+    LLVMIRGenerationRefCount -= 1;
+    if (LLVMIRGenerationRefCount == 0)
+      LLVMIRGeneration.stopTimer();
+  }
 
-    void HandleInlineFunctionDefinition(FunctionDecl *D) override {
-      PrettyStackTraceDecl CrashInfo(D, SourceLocation(),
-                                     Context->getSourceManager(),
-                                     "LLVM IR generation of inline function");
-      if (TimerIsEnabled)
-        LLVMIRGeneration.startTimer();
+  return true;
+}
 
-      Gen->HandleInlineFunctionDefinition(D);
+void BackendConsumer::HandleInlineFunctionDefinition(FunctionDecl *D) {
+  PrettyStackTraceDecl CrashInfo(D, SourceLocation(),
+                                 Context->getSourceManager(),
+                                 "LLVM IR generation of inline function");
+  if (TimerIsEnabled)
+    LLVMIRGeneration.startTimer();
 
-      if (TimerIsEnabled)
-        LLVMIRGeneration.stopTimer();
-    }
+  Gen->HandleInlineFunctionDefinition(D);
 
-    void HandleInterestingDecl(DeclGroupRef D) override {
-      // Ignore interesting decls from the AST reader after IRGen is finished.
-      if (!IRGenFinished)
-        HandleTopLevelDecl(D);
-    }
+  if (TimerIsEnabled)
+    LLVMIRGeneration.stopTimer();
+}
 
-    // Links each entry in LinkModules into our module.  Returns true on error.
-    bool LinkInModules(llvm::Module *M) {
-      for (auto &LM : LinkModules) {
-        assert(LM.Module && "LinkModule does not actually have a module");
-        if (LM.PropagateAttrs)
-          for (Function &F : *LM.Module) {
-            // Skip intrinsics. Keep consistent with how intrinsics are created
-            // in LLVM IR.
-            if (F.isIntrinsic())
-              continue;
-            CodeGen::mergeDefaultFunctionDefinitionAttributes(
-                F, CodeGenOpts, LangOpts, TargetOpts, LM.Internalize);
-          }
-
-        CurLinkModule = LM.Module.get();
-
-        bool Err;
-        if (LM.Internalize) {
-          Err = Linker::linkModules(
-              *M, std::move(LM.Module), LM.LinkFlags,
-              [](llvm::Module &M, const llvm::StringSet<> &GVS) {
-                internalizeModule(M, [&GVS](const llvm::GlobalValue &GV) {
-                  return !GV.hasName() || (GVS.count(GV.getName()) == 0);
-                });
-              });
-        } else {
-          Err = Linker::linkModules(*M, std::move(LM.Module), LM.LinkFlags);
-        }
+void BackendConsumer::HandleInterestingDecl(DeclGroupRef D) {
+  // Ignore interesting decls from the AST reader after IRGen is finished.
+  if (!IRGenFinished)
+    HandleTopLevelDecl(D);
+}
 
-        if (Err)
-          return true;
+// Links each entry in LinkModules into our module.  Returns true on error.
+bool BackendConsumer::LinkInModules(llvm::Module *M, bool ShouldLinkFiles) {
+
+  for (auto &LM : LinkModules) {
+    assert(LM.Module && "LinkModule does not actually have a module");
+
+    // If ShouldLinkFiles is not set, skip files added via the
+    // -mlink-bitcode-files, only linking -mlink-builtin-bitcode
+    if (!LM.Internalize && !ShouldLinkFiles)
+      continue;
+
+    if (LM.PropagateAttrs)
+      for (Function &F : *LM.Module) {
+        // Skip intrinsics. Keep consistent with how intrinsics are created
+        // in LLVM IR.
+        if (F.isIntrinsic())
+          continue;
+        CodeGen::mergeDefaultFunctionDefinitionAttributes(
+          F, CodeGenOpts, LangOpts, TargetOpts, LM.Internalize);
       }
-      LinkModules.clear();
-      return false; // success
-    }
 
-    void HandleTranslationUnit(ASTContext &C) override {
-      {
-        llvm::TimeTraceScope TimeScope("Frontend");
-        PrettyStackTraceString CrashInfo("Per-file LLVM IR generation");
-        if (TimerIsEnabled) {
-          LLVMIRGenerationRefCount += 1;
-          if (LLVMIRGenerationRefCount == 1)
-            LLVMIRGeneration.startTimer();
-        }
+    CurLinkModule = LM.Module.get();
+
+    // TODO: If CloneModule() is updated to support cloning of unmaterialized
+    // modules, we can remove this
+    bool Err;
+    if (Error E = CurLinkModule->materializeAll())
+      return false;
+
+    // Create a Clone to move to the linker, which preserves the original
+    // linking modules, allowing them to be linked again in the future
+    // TODO: Add a ShouldCleanup option to make Cloning optional. When
+    // set, we can pass the original modules to the linker for cleanup
+    std::unique_ptr<llvm::Module> Clone = llvm::CloneModule(*LM.Module);
+
+    if (LM.Internalize) {
+      Err = Linker::linkModules(
+          *M, std::move(Clone), LM.LinkFlags,
+          [](llvm::Module &M, const llvm::StringSet<> &GVS) {
+            internalizeModule(M, [&GVS](const llvm::GlobalValue &GV) {
+              return !GV.hasName() || (GVS.count(GV.getName()) == 0);
+            });
+          });
+    } else
+      Err = Linker::linkModules(*M, std::move(Clone), LM.LinkFlags);
+
+    if (Err)
+      return true;
+  }
 
-        Gen->HandleTranslationUnit(C);
+  return false; // success
+}
 
-        if (TimerIsEnabled) {
-          LLVMIRGenerationRefCount -= 1;
-          if (LLVMIRGenerationRefCount == 0)
-            LLVMIRGeneration.stopTimer();
-        }
+void BackendConsumer::HandleTranslationUnit(ASTContext &C) {
+  {
+    llvm::TimeTraceScope TimeScope("Frontend");
+    PrettyStackTraceString CrashInfo("Per-file LLVM IR generation");
+    if (TimerIsEnabled) {
+      LLVMIRGenerationRefCount += 1;
+      if (LLVMIRGenerationRefCount == 1)
+        LLVMIRGeneration.startTimer();
+    }
 
-        IRGenFinished = true;
-      }
+    Gen->HandleTranslationUnit(C);
 
-      // Silently ignore if we weren't initialized for some reason.
-      if (!getModule())
-        return;
-
-      LLVMContext &Ctx = getModule()->getContext();
-      std::unique_ptr<DiagnosticHandler> OldDiagnosticHandler =
-          Ctx.getDiagnosticHandler();
-      Ctx.setDiagnosticHandler(std::make_unique<ClangDiagnosticHandler>(
-        CodeGenOpts, this));
-
-      Expected<std::unique_ptr<llvm::ToolOutputFile>> OptRecordFileOrErr =
-          setupLLVMOptimizationRemarks(
-              Ctx, CodeGenOpts.OptRecordFile, CodeGenOpts.OptRecordPasses,
-              CodeGenOpts.OptRecordFormat, CodeGenOpts.DiagnosticsWithHotness,
-              CodeGenOpts.DiagnosticsHotnessThreshold);
-
-      if (Error E = OptRecordFileOrErr.takeError()) {
-        reportOptRecordError(std::move(E), Diags, CodeGenOpts);
-        return;
-      }
+    if (TimerIsEnabled) {
+      LLVMIRGenerationRefCount -= 1;
+      if (LLVMIRGenerationRefCount == 0)
+        LLVMIRGeneration.stopTimer();
+    }
 
-      std::unique_ptr<llvm::ToolOutputFile> OptRecordFile =
-          std::move(*OptRecordFileOrErr);
+    IRGenFinished = true;
+  }
 
-      if (OptRecordFile &&
-          CodeGenOpts.getProfileUse() != CodeGenOptions::ProfileNone)
-        Ctx.setDiagnosticsHotnessRequested(true);
+  // Silently ignore if we weren't initialized for some reason.
+  if (!getModule())
+    return;
 
-      if (CodeGenOpts.MisExpect) {
-        Ctx.setMisExpectWarningRequested(true);
-      }
+  LLVMContext &Ctx = getModule()->getContext();
+  std::unique_ptr<DiagnosticHandler> OldDiagnosticHandler =
+    Ctx.getDiagnosticHandler();
+  Ctx.setDiagnosticHandler(std::make_unique<ClangDiagnosticHandler>(
+      CodeGenOpts, this));
 
-      if (CodeGenOpts.DiagnosticsMisExpectTolerance) {
-        Ctx.setDiagnosticsMisExpectTolerance(
-            CodeGenOpts.DiagnosticsMisExpectTolerance);
-      }
+  Expected<std::unique_ptr<llvm::ToolOutputFile>> OptRecordFileOrErr =
+    setupLLVMOptimizationRemarks(
+      Ctx, CodeGenOpts.OptRecordFile, CodeGenOpts.OptRecordPasses,
+      CodeGenOpts.OptRecordFormat, CodeGenOpts.DiagnosticsWithHotness,
+      CodeGenOpts.DiagnosticsHotnessThreshold);
 
-      // Link each LinkModule into our module.
-      if (LinkInModules(getModule()))
-        return;
+  if (Error E = OptRecordFileOrErr.takeError()) {
+    reportOptRecordError(std::move(E), Diags, CodeGenOpts);
+    return;
+  }
 
-      for (auto &F : getModule()->functions()) {
-        if (const Decl *FD = Gen->GetDeclForMangledName(F.getName())) {
-          auto Loc = FD->getASTContext().getFullLoc(FD->getLocation());
-          // TODO: use a fast content hash when available.
-          auto NameHash = llvm::hash_value(F.getName());
-          ManglingFullSourceLocs.push_back(std::make_pair(NameHash, Loc));
-        }
-      }
+  std::unique_ptr<llvm::ToolOutputFile> OptRecordFile =
+    std::move(*OptRecordFileOrErr);
 
-      if (CodeGenOpts.ClearASTBeforeBackend) {
-        LLVM_DEBUG(llvm::dbgs() << "Clearing AST...\n");
-        // Access to the AST is no longer available after this.
-        // Other things that the ASTContext manages are still available, e.g.
-        // the SourceManager. It'd be nice if we could separate out all the
-        // things in ASTContext used after this point and null out the
-        // ASTContext, but too many various parts of the ASTContext are still
-        // used in various parts.
-        C.cleanup();
-        C.getAllocator().Reset();
-      }
+  if (OptRecordFile &&
+      CodeGenOpts.getProfileUse() != CodeGenOptions::ProfileNone)
+    Ctx.setDiagnosticsHotnessRequested(true);
 
-      EmbedBitcode(getModule(), CodeGenOpts, llvm::MemoryBufferRef());
+  if (CodeGenOpts.MisExpect) {
+    Ctx.setMisExpectWarningRequested(true);
+  }
 
-      EmitBackendOutput(Diags, HeaderSearchOpts, CodeGenOpts, TargetOpts,
-                        LangOpts, C.getTargetInfo().getDataLayoutString(),
-                        getModule(), Action, FS, std::move(AsmOutStream));
+  if (CodeGenOpts.DiagnosticsMisExpectTolerance) {
+    Ctx.setDiagnosticsMisExpectTolerance(
+      CodeGenOpts.DiagnosticsMisExpectTolerance);
+  }
 
-      Ctx.setDiagnosticHandler(std::move(OldDiagnosticHandler));
+  // Link each LinkModule into our module.
+  if (LinkInModules(getModule()))
+    return;
 
-      if (OptRecordFile)
-        OptRecordFile->keep();
+  for (auto &F : getModule()->functions()) {
+    if (const Decl *FD = Gen->GetDeclForMangledName(F.getName())) {
+      auto Loc = FD->getASTContext().getFullLoc(FD->getLocation());
+      // TODO: use a fast content hash when available.
+      auto NameHash = llvm::hash_value(F.getName());
+      ManglingFullSourceLocs.push_back(std::make_pair(NameHash, Loc));
     }
+  }
 
-    void HandleTagDeclDefinition(TagDecl *D) override {
-      PrettyStackTraceDecl CrashInfo(D, SourceLocation(),
-                                     Context->getSourceManager(),
-                                     "LLVM IR generation of declaration");
-      Gen->HandleTagDeclDefinition(D);
-    }
+  if (CodeGenOpts.ClearASTBeforeBackend) {
+    LLVM_DEBUG(llvm::dbgs() << "Clearing AST...\n");
+    // Access to the AST is no longer available after this.
+    // Other things that the ASTContext manages are still available, e.g.
+    // the SourceManager. It'd be nice if we could separate out all the
+    // things in ASTContext used after this point and null out the
+    // ASTContext, but too many various parts of the ASTContext are still
+    // used in various parts.
+    C.cleanup();
+    C.getAllocator().Reset();
+  }
 
-    void HandleTagDeclRequiredDefinition(const TagDecl *D) override {
-      Gen->HandleTagDeclRequiredDefinition(D);
-    }
+  EmbedBitcode(getModule(), CodeGenOpts, llvm::MemoryBufferRef());
 
-    void CompleteTentativeDefinition(VarDecl *D) override {
-      Gen->CompleteTentativeDefinition(D);
-    }
+  EmitBackendOutput(Diags, HeaderSearchOpts, CodeGenOpts, TargetOpts, LangOpts,
+                    C.getTargetInfo().getDataLayoutString(), getModule(),
+                    Action, FS, std::move(AsmOutStream), this);
 
-    void CompleteExternalDeclaration(VarDecl *D) override {
-      Gen->CompleteExternalDeclaration(D);
-    }
+  Ctx.setDiagnosticHandler(std::move(OldDiagnosticHandler));
 
-    void AssignInheritanceModel(CXXRecordDecl *RD) override {
-      Gen->AssignInheritanceModel(RD);
-    }
+  if (OptRecordFile)
+    OptRecordFile->keep();
+}
 
-    void HandleVTable(CXXRecordDecl *RD) override {
-      Gen->HandleVTable(RD);
-    }
+void BackendConsumer::HandleTagDeclDefinition(TagDecl *D) {
+  PrettyStackTraceDecl CrashInfo(D, SourceLocation(),
+                                 Context->getSourceManager(),
+                                 "LLVM IR generation of declaration");
+  Gen->HandleTagDeclDefinition(D);
+}
 
-    /// Get the best possible source location to represent a diagnostic that
-    /// may have associated debug info.
-    const FullSourceLoc
-    getBestLocationFromDebugLoc(const llvm::DiagnosticInfoWithLocationBase &D,
-                                bool &BadDebugInfo, StringRef &Filename,
-                                unsigned &Line, unsigned &Column) const;
-
-    std::optional<FullSourceLoc>
-    getFunctionSourceLocation(const Function &F) const;
-
-    void DiagnosticHandlerImpl(const llvm::DiagnosticInfo &DI);
-    /// Specialized handler for InlineAsm diagnostic.
-    /// \return True if the diagnostic has been successfully reported, false
-    /// otherwise.
-    bool InlineAsmDiagHandler(const llvm::DiagnosticInfoInlineAsm &D);
-    /// Specialized handler for diagnostics reported using SMDiagnostic.
-    void SrcMgrDiagHandler(const llvm::DiagnosticInfoSrcMgr &D);
-    /// Specialized handler for StackSize diagnostic.
-    /// \return True if the diagnostic has been successfully reported, false
-    /// otherwise.
-    bool StackSizeDiagHandler(const llvm::DiagnosticInfoStackSize &D);
-    /// Specialized handler for ResourceLimit diagnostic.
-    /// \return True if the diagnostic has been successfully reported, false
-    /// otherwise.
-    bool ResourceLimitDiagHandler(const llvm::DiagnosticInfoResourceLimit &D);
-
-    /// Specialized handler for unsupported backend feature diagnostic.
-    void UnsupportedDiagHandler(const llvm::DiagnosticInfoUnsupported &D);
-    /// Specialized handlers for optimization remarks.
-    /// Note that these handlers only accept remarks and they always handle
-    /// them.
-    void EmitOptimizationMessage(const llvm::DiagnosticInfoOptimizationBase &D,
-                                 unsigned DiagID);
-    void
-    OptimizationRemarkHandler(const llvm::DiagnosticInfoOptimizationBase &D);
-    void OptimizationRemarkHandler(
-        const llvm::OptimizationRemarkAnalysisFPCommute &D);
-    void OptimizationRemarkHandler(
-        const llvm::OptimizationRemarkAnalysisAliasing &D);
-    void OptimizationFailureHandler(
-        const llvm::DiagnosticInfoOptimizationFailure &D);
-    void DontCallDiagHandler(const DiagnosticInfoDontCall &D);
-    /// Specialized handler for misexpect warnings.
-    /// Note that misexpect remarks are emitted through ORE
-    void MisExpectDiagHandler(const llvm::DiagnosticInfoMisExpect &D);
-  };
+void BackendConsumer::HandleTagDeclRequiredDefinition(const TagDecl *D) {
+  Gen->HandleTagDeclRequiredDefinition(D);
+}
 
-  void BackendConsumer::anchor() {}
+void BackendConsumer::CompleteTentativeDefinition(VarDecl *D) {
+  Gen->CompleteTentativeDefinition(D);
 }
 
+void BackendConsumer::CompleteExternalDeclaration(VarDecl *D) {
+  Gen->CompleteExternalDeclaration(D);
+}
+
+void BackendConsumer::AssignInheritanceModel(CXXRecordDecl *RD) {
+  Gen->AssignInheritanceModel(RD);
+}
+
+void BackendConsumer::HandleVTable(CXXRecordDecl *RD) {
+  Gen->HandleVTable(RD);
+}
+
+void BackendConsumer::anchor() { }
+
+} // namespace clang
+
 bool ClangDiagnosticHandler::handleDiagnostics(const DiagnosticInfo &DI) {
   BackendCon->DiagnosticHandlerImpl(DI);
   return true;