[CodeGen] Implement post-optimization linking option for builtin bitcodes

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. 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 defintion.

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

Author: lamb-j

clang/include/clang/CodeGen/BackendUtil.h

@@ -30,6 +30,7 @@ namespace clang {
   class CodeGenOptions;
   class TargetOptions;
   class LangOptions;
+  class BackendConsumer;
 
   enum BackendAction {
     Backend_EmitAssembly,  ///< Emit native assembly files
@@ -45,7 +46,8 @@ namespace clang {
                          const TargetOptions &TOpts, const LangOptions &LOpts,
                          StringRef TDesc, llvm::Module *M, BackendAction Action,
                          llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
-                         std::unique_ptr<raw_pwrite_stream> OS);
+                         std::unique_ptr<raw_pwrite_stream> OS,
+                         BackendConsumer *BC = nullptr);
 
   void EmbedBitcode(llvm::Module *M, const CodeGenOptions &CGOpts,
                     llvm::MemoryBufferRef Buf);

clang/lib/CodeGen/BackendConsumer.h

@@ -113,7 +113,7 @@ class BackendConsumer : public ASTConsumer {
 
 
  // Links each entry in LinkModules into our module.  Returns true on error.
-  bool LinkInModules(llvm::Module *M);
+  bool LinkInModules(llvm::Module *M, bool ShouldLinkFiles = true);
 
   /// Get the best possible source location to represent a diagnostic that
   /// may have associated debug info.

clang/lib/CodeGen/BackendUtil.cpp

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/CodeGen/BackendUtil.h"
+#include "BackendConsumer.h"
+#include "LinkInModulesPass.h"
 #include "clang/Basic/CodeGenOptions.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/LangOptions.h"
@@ -98,6 +100,11 @@ extern cl::opt<bool> PrintPipelinePasses;
 static cl::opt<bool> ClSanitizeOnOptimizerEarlyEP(
     "sanitizer-early-opt-ep", cl::Optional,
     cl::desc("Insert sanitizers on OptimizerEarlyEP."), cl::init(false));
+
+// Re-link builtin bitcodes after optimization
+static cl::opt<bool> ClRelinkBuiltinBitcodePostop(
+    "relink-builtin-bitcode-postop", cl::Optional,
+    cl::desc("Re-link builtin bitcodes after optimization."), cl::init(false));
 }
 
 namespace {
@@ -156,10 +163,9 @@ class EmitAssemblyHelper {
     return F;
   }
 
-  void
-  RunOptimizationPipeline(BackendAction Action,
-                          std::unique_ptr<raw_pwrite_stream> &OS,
-                          std::unique_ptr<llvm::ToolOutputFile> &ThinLinkOS);
+  void RunOptimizationPipeline(
+      BackendAction Action, std::unique_ptr<raw_pwrite_stream> &OS,
+      std::unique_ptr<llvm::ToolOutputFile> &ThinLinkOS, BackendConsumer *BC);
   void RunCodegenPipeline(BackendAction Action,
                           std::unique_ptr<raw_pwrite_stream> &OS,
                           std::unique_ptr<llvm::ToolOutputFile> &DwoOS);
@@ -194,8 +200,8 @@ class EmitAssemblyHelper {
   std::unique_ptr<TargetMachine> TM;
 
   // Emit output using the new pass manager for the optimization pipeline.
-  void EmitAssembly(BackendAction Action,
-                    std::unique_ptr<raw_pwrite_stream> OS);
+  void EmitAssembly(BackendAction Action, std::unique_ptr<raw_pwrite_stream> OS,
+                    BackendConsumer *BC);
 };
 }
 
@@ -764,7 +770,7 @@ static void addSanitizers(const Triple &TargetTriple,
 
 void EmitAssemblyHelper::RunOptimizationPipeline(
     BackendAction Action, std::unique_ptr<raw_pwrite_stream> &OS,
-    std::unique_ptr<llvm::ToolOutputFile> &ThinLinkOS) {
+    std::unique_ptr<llvm::ToolOutputFile> &ThinLinkOS, BackendConsumer *BC) {
   std::optional<PGOOptions> PGOOpt;
 
   if (CodeGenOpts.hasProfileIRInstr())
@@ -1040,6 +1046,13 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
     }
   }
 
+  // Re-link against any bitcodes supplied via the -mlink-builtin-bitcode option
+  // Some optimizations may generate new function calls that would not have
+  // been linked pre-optimization (i.e. fused sincos calls generated by
+  // AMDGPULibCalls::fold_sincos.)
+  if (ClRelinkBuiltinBitcodePostop)
+    MPM.addPass(LinkInModulesPass(BC, false));
+
   // Add a verifier pass if requested. We don't have to do this if the action
   // requires code generation because there will already be a verifier pass in
   // the code-generation pipeline.
@@ -1169,7 +1182,8 @@ void EmitAssemblyHelper::RunCodegenPipeline(
 }
 
 void EmitAssemblyHelper::EmitAssembly(BackendAction Action,
-                                      std::unique_ptr<raw_pwrite_stream> OS) {
+                                      std::unique_ptr<raw_pwrite_stream> OS,
+                                      BackendConsumer *BC) {
   TimeRegion Region(CodeGenOpts.TimePasses ? &CodeGenerationTime : nullptr);
   setCommandLineOpts(CodeGenOpts);
 
@@ -1185,7 +1199,7 @@ void EmitAssemblyHelper::EmitAssembly(BackendAction Action,
   cl::PrintOptionValues();
 
   std::unique_ptr<llvm::ToolOutputFile> ThinLinkOS, DwoOS;
-  RunOptimizationPipeline(Action, OS, ThinLinkOS);
+  RunOptimizationPipeline(Action, OS, ThinLinkOS, BC);
   RunCodegenPipeline(Action, OS, DwoOS);
 
   if (ThinLinkOS)
@@ -1299,14 +1313,12 @@ static void runThinLTOBackend(
   }
 }
 
-void clang::EmitBackendOutput(DiagnosticsEngine &Diags,
-                              const HeaderSearchOptions &HeaderOpts,
-                              const CodeGenOptions &CGOpts,
-                              const clang::TargetOptions &TOpts,
-                              const LangOptions &LOpts, StringRef TDesc,
-                              llvm::Module *M, BackendAction Action,
-                              IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
-                              std::unique_ptr<raw_pwrite_stream> OS) {
+void clang::EmitBackendOutput(
+    DiagnosticsEngine &Diags, const HeaderSearchOptions &HeaderOpts,
+    const CodeGenOptions &CGOpts, const clang::TargetOptions &TOpts,
+    const LangOptions &LOpts, StringRef TDesc, llvm::Module *M,
+    BackendAction Action, IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+    std::unique_ptr<raw_pwrite_stream> OS, BackendConsumer *BC) {
 
   llvm::TimeTraceScope TimeScope("Backend");
 
@@ -1349,7 +1361,7 @@ void clang::EmitBackendOutput(DiagnosticsEngine &Diags,
   }
 
   EmitAssemblyHelper AsmHelper(Diags, HeaderOpts, CGOpts, TOpts, LOpts, M, VFS);
-  AsmHelper.EmitAssembly(Action, std::move(OS));
+  AsmHelper.EmitAssembly(Action, std::move(OS), BC);
 
   // Verify clang's TargetInfo DataLayout against the LLVM TargetMachine's
   // DataLayout.

clang/lib/CodeGen/CMakeLists.txt

@@ -103,6 +103,7 @@ add_clang_library(clangCodeGen
   ConstantInitBuilder.cpp
   CoverageMappingGen.cpp
   ItaniumCXXABI.cpp
+  LinkInModulesPass.cpp
   MacroPPCallbacks.cpp
   MicrosoftCXXABI.cpp
   ModuleBuilder.cpp

clang/lib/CodeGen/CodeGenAction.cpp

@@ -49,6 +49,7 @@
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Transforms/IPO/Internalize.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 
 #include <optional>
 using namespace clang;
@@ -229,9 +230,16 @@ void BackendConsumer::HandleInterestingDecl(DeclGroupRef D) {
 }
 
 // Links each entry in LinkModules into our module.  Returns true on error.
-bool BackendConsumer::LinkInModules(llvm::Module *M) {
+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
@@ -244,24 +252,33 @@ bool BackendConsumer::LinkInModules(llvm::Module *M) {
 
     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(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);
-    }
+          *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;
   }
-  LinkModules.clear();
+
   return false; // success
 }
 
@@ -350,9 +367,9 @@ void BackendConsumer::HandleTranslationUnit(ASTContext &C) {
 
   EmbedBitcode(getModule(), CodeGenOpts, llvm::MemoryBufferRef());
 
-  EmitBackendOutput(Diags, HeaderSearchOpts, CodeGenOpts, TargetOpts,
-                    LangOpts, C.getTargetInfo().getDataLayoutString(),
-                    getModule(), Action, FS, std::move(AsmOutStream));
+  EmitBackendOutput(Diags, HeaderSearchOpts, CodeGenOpts, TargetOpts, LangOpts,
+                    C.getTargetInfo().getDataLayoutString(), getModule(),
+                    Action, FS, std::move(AsmOutStream), this);
 
   Ctx.setDiagnosticHandler(std::move(OldDiagnosticHandler));
 

clang/lib/CodeGen/LinkInModulesPass.cpp

@@ -0,0 +1,29 @@
+//===-- LinkInModulesPass.cpp - Module Linking pass --------------- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// LinkInModulesPass implementation.
+///
+//===----------------------------------------------------------------------===//
+
+#include "LinkInModulesPass.h"
+#include "BackendConsumer.h"
+
+using namespace llvm;
+
+LinkInModulesPass::LinkInModulesPass(clang::BackendConsumer *BC,
+                                     bool ShouldLinkFiles)
+    : BC(BC), ShouldLinkFiles(ShouldLinkFiles) {}
+
+PreservedAnalyses LinkInModulesPass::run(Module &M, ModuleAnalysisManager &AM) {
+
+  if (BC && BC->LinkInModules(&M, ShouldLinkFiles))
+    report_fatal_error("Bitcode module linking failed, compilation aborted!");
+
+  return PreservedAnalyses::all();
+}

clang/lib/CodeGen/LinkInModulesPass.h

@@ -0,0 +1,42 @@
+//===-- LinkInModulesPass.h - Module Linking pass ----------------- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+///
+/// This file provides a pass to link in Modules from a provided
+/// BackendConsumer.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_BITCODE_LINKINMODULESPASS_H
+#define LLVM_BITCODE_LINKINMODULESPASS_H
+
+#include "BackendConsumer.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+class Module;
+class ModulePass;
+class Pass;
+
+/// Create and return a pass that links in Moduels from a provided
+/// BackendConsumer to a given primary Module. Note that this pass is designed
+/// for use with the legacy pass manager.
+class LinkInModulesPass : public PassInfoMixin<LinkInModulesPass> {
+  clang::BackendConsumer *BC;
+  bool ShouldLinkFiles;
+
+public:
+  LinkInModulesPass(clang::BackendConsumer *BC, bool ShouldLinkFiles = true);
+
+  PreservedAnalyses run(Module &M, AnalysisManager<Module> &);
+  static bool isRequired() { return true; }
+};
+
+} // namespace llvm
+
+#endif