[LTO] Add a config option to make the final output bitcode

[jhuber6]

Summary

The LTO pass currently only offers assembly or object file outputs. We
can obtain the final LLVM output using the hooks, but is inconvenient
and disrupts the flow of the compilation. This patch proposes adding a
convenience function to perform this if the user requires it. We would
like this for the desire to retain the post-optimized LTO bitcode for
LLVM offloading compileations, such as JIT, builtin linking, or
performing a relocatable link on LLVM-IR.

Depends on D152442

[jhuber6]

Remove leftover debug

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

std::vector<std::string> getTargetFeatures(ArrayRef<OffloadFile> InputFiles) {
   for (StringRef Feature : llvm::reverse(Features)) {
     if (UsedFeatures.insert(Feature.drop_front()).second)
       UnifiedFeatures.push_back(Feature.str());
   }
 
   return UnifiedFeatures;
 }
 
-template <typename ModuleHook = function_ref<bool(size_t, const Module &)>>
-std::unique_ptr<lto::LTO> createLTO(
-    const ArgList &Args, const std::vector<std::string> &Features,
-    ModuleHook Hook = [](size_t, const Module &) { return true; }) {
+std::unique_ptr<lto::LTO> createLTO(const ArgList &Args,
+                                    const std::vector<std::string> &Features) {
   const llvm::Triple Triple(Args.getLastArgValue(OPT_triple_EQ));
   StringRef Arch = Args.getLastArgValue(OPT_arch_EQ);
   lto::Config Conf;
   lto::ThinBackend Backend;
   // TODO: Handle index-only thin-LTO
   Backend =
       lto::createInProcessThinBackend(llvm::heavyweight_hardware_concurrency());
 
Conf.PreCodeGenModuleHook = [=](size_t Task, const Module &M) {
       error_code EC;
       raw_fd_ostream LinkedBitcode(File, EC, sys::fs::OF_None);
       if (EC)
         reportError(errorCodeToError(EC));
       WriteBitcodeToFile(M, LinkedBitcode);
       return true;
     };
   }
-  Conf.PostOptModuleHook = Hook;
+  Conf.OutputPreCodegenBitcode = Args.hasArg(OPT_embed_bitcode) ||
+                                 Args.hasArg(OPT_builtin_bitcode_EQ) ||
+                                 Args.hasArg(OPT_clang_backend);
   Conf.CGFileType =
       (Triple.isNVPTX() || SaveTemps) ? CGFT_AssemblyFile : CGFT_ObjectFile;
 
   // TODO: Handle remark files
   Conf.HasWholeProgramVisibility = Args.hasArg(OPT_whole_program);
 
   return std::make_unique<lto::LTO>(std::move(Conf), Backend);
 }
Error linkBitcodeFiles(SmallVectorImpl<OffloadFile> &InputFiles,
   }
 
   if (BitcodeInputFiles.empty())
     return Error::success();
 
   // Remove all the bitcode files that we moved from the original input.
   llvm::erase_if(InputFiles, [](OffloadFile &F) { return !F.getBinary(); });
 
-  // LTO Module hook to output bitcode without running the backend.
-  SmallVector<StringRef> BitcodeOutput;
-  auto OutputBitcode = [&](size_t, const Module &M) {
-    auto TempFileOrErr = createOutputFile(sys::path::filename(ExecutableName) +
-                                              "-jit-" + Triple.getTriple(),
-                                          "bc");
-    if (!TempFileOrErr)
-      reportError(TempFileOrErr.takeError());
-
-    std::error_code EC;
-    raw_fd_ostream LinkedBitcode(*TempFileOrErr, EC, sys::fs::OF_None);
-    if (EC)
-      reportError(errorCodeToError(EC));
-    WriteBitcodeToFile(M, LinkedBitcode);
-    BitcodeOutput.push_back(*TempFileOrErr);
-    return false;
-  };
-
   // We assume visibility of the whole program if every input file was bitcode.
   auto Features = getTargetFeatures(BitcodeInputFiles);
-  auto LTOBackend = Args.hasArg(OPT_embed_bitcode) ||
-                            Args.hasArg(OPT_builtin_bitcode_EQ) ||
-                            Args.hasArg(OPT_clang_backend)
-                        ? createLTO(Args, Features, OutputBitcode)
-                        : createLTO(Args, Features);
+  auto LTOBackend = createLTO(Args, Features);
 
   // We need to resolve the symbols so the LTO backend knows which symbols need
   // to be kept or can be internalized. This is a simplified symbol resolution
   // scheme to approximate the full resolution a linker would do.
   uint64_t Idx = 0;
   DenseSet<StringRef> PrevailingSymbols;
   for (auto &BitcodeInput : BitcodeInputFiles) {
     // Get a semi-unique buffer identifier for Thin-LTO.
for (auto &BitcodeInput : BitcodeInputFiles) {
     // Add the bitcode file with its resolved symbols to the LTO job.
     if (Error Err = LTOBackend->add(std::move(*BitcodeFileOrErr), Resolutions))
       return Err;
   }
 
   // Run the LTO job to compile the bitcode.
   size_t MaxTasks = LTOBackend->getMaxTasks();
   SmallVector<StringRef> Files(MaxTasks);
-  auto AddStream =
-      [&](size_t Task,
-          const Twine &ModuleName) -> std::unique_ptr<CachedFileStream> {
+  auto AddStream = [&](size_t Task, const Twine &ModuleName) {
+    bool OutputBitcode = Args.hasArg(OPT_embed_bitcode) ||
+                         Args.hasArg(OPT_builtin_bitcode_EQ) ||
+                         Args.hasArg(OPT_clang_backend);
     int FD = -1;
     auto &TempFile = Files[Task];
-    StringRef Extension = (Triple.isNVPTX() || SaveTemps) ? "s" : "o";
+    StringRef Extension = OutputBitcode                     ? "bc"
+                          : (Triple.isNVPTX() || SaveTemps) ? "s"
+                                                            : "o";
     std::string TaskStr = Task ? "." + std::to_string(Task) : "";
     auto TempFileOrErr =
         createOutputFile(sys::path::filename(ExecutableName) + "." +
                              Triple.getTriple() + "." + Arch + TaskStr,
                          Extension);
     if (!TempFileOrErr)
       reportError(TempFileOrErr.takeError());
     TempFile = *TempFileOrErr;
     if (std::error_code EC = sys::fs::openFileForWrite(TempFile, FD))
       reportError(errorCodeToError(EC));
     return std::make_unique<CachedFileStream>(
         std::make_unique<llvm::raw_fd_ostream>(FD, true));
   };
 
   if (Error Err = LTOBackend->run(AddStream))
     return Err;
 
   if (LTOError)
     return createStringError(inconvertibleErrorCode(),
                              "Errors encountered inside the LTO pipeline.");
 
-  // If we are embedding bitcode we only need the intermediate output.
-  bool SingleOutput = Files.size() == 1;
-  if (Args.hasArg(OPT_embed_bitcode)) {
-    if (BitcodeOutput.size() != 1 || !SingleOutput)
+  if (Args.hasArg(OPT_embed_bitcode))
+    if (Files.size() != 1)
       return createStringError(inconvertibleErrorCode(),
                                "Cannot embed bitcode with multiple files.");
-    OutputFiles.push_back(Args.MakeArgString(BitcodeOutput.front()));
-    return Error::success();
-  }
 
-  // Append the new inputs to the device linker input. If the user requested an
-  // internalizing link we need to pass the bitcode to clang.
-  for (StringRef File :
-       Args.hasArg(OPT_clang_backend) || Args.hasArg(OPT_builtin_bitcode_EQ)
-           ? BitcodeOutput
-           : Files)
+  // Append the new inputs to the device linker input.
+  for (StringRef File : Files)
     OutputFiles.push_back(File);
 
   return Error::success();
 }
 
 Expected<StringRef> writeOffloadFile(const OffloadFile &File) {
   const OffloadBinary &Binary = *File.getBinary();
 

llvm/include/llvm/LTO/Config.h

struct Config {
   /// link.
   bool HasWholeProgramVisibility = false;
 
   /// Always emit a Regular LTO object even when it is empty because no Regular
   /// LTO modules were linked. This option is useful for some build system which
   /// want to know a priori all possible output files.
   bool AlwaysEmitRegularLTOObj = false;
 
+  /// Write the pre-codegen module to the output without performing codegen.
+  /// Useful for post-processing on the fully linked LTO module.
+  bool OutputPreCodegenBitcode = false;
+
   /// Allows non-imported definitions to get the potentially more constraining
   /// visibility from the prevailing definition. FromPrevailing is the default
   /// because it works for many binary formats. ELF can use the more optimized
   /// 'ELF' scheme.
   VisScheme VisibilityScheme = FromPrevailing;
 
   /// If this field is set, the set of passes run in the middle-end optimizer
   /// will be the one specified by the string. Only works with the new pass

llvm/lib/LTO/LTOBackend.cpp

 }
 
 static void codegen(const Config &Conf, TargetMachine *TM,
                     AddStreamFn AddStream, unsigned Task, Module &Mod,
                     const ModuleSummaryIndex &CombinedIndex) {
   if (Conf.PreCodeGenModuleHook && !Conf.PreCodeGenModuleHook(Task, Mod))
     return;
 
+  if (Conf.OutputPreCodegenBitcode) {
+    Expected<std::unique_ptr<CachedFileStream>> StreamOrErr =
+        AddStream(Task, Mod.getModuleIdentifier());
+    if (Error Err = StreamOrErr.takeError())
+      report_fatal_error(std::move(Err));
+    std::unique_ptr<CachedFileStream> &Stream = *StreamOrErr;
+
+    WriteBitcodeToFile(Mod, *Stream->OS);
+    return;
+  }
+
   if (EmbedBitcode == LTOBitcodeEmbedding::EmbedOptimized)
     llvm::embedBitcodeInModule(Mod, llvm::MemoryBufferRef(),
                                /*EmbedBitcode*/ true,
                                /*EmbedCmdline*/ false,
                                /*CmdArgs*/ std::vector<uint8_t>());
 
   std::unique_ptr<ToolOutputFile> DwoOut;
   SmallString<1024> DwoFile(Conf.SplitDwarfOutput);