[OpenMP] Add support for embedding bitcode images in wrapper tool

[jhuber6]

Summary

Summary;
This patch adds support for embedding device images in the linker
wrapper tool. This will be used for performing JIT functionality in the
future.

Depends on D117048

[jhuber6]

Update approach, improve LTO file handling for multiple outputs and don't change the LTO config.

[jhuber6]

Forgot to use hook.

[saiislam]

It seems that this patch along with D117156 and D117246 is giving patch application failed error [https://buildkite.com/llvm-project/diff-checks/builds/82688].
arc patch is also giving the same error.

[jhuber6]

In D117049#3255748, @saiislam wrote:

It seems that this patch along with D117156 and D117246 is giving patch application failed error [https://buildkite.com/llvm-project/diff-checks/builds/82688].
arc patch is also giving the same error.

I think what happened is I didn't include a small fix patch I had locally in the list somewhere which made the diffs not apply. I could probably try to squash them and reapply if it's important.

[jhuber6]

Rework commits.

[jhuber6]

Removing clang flag because LTO won't be supported when these land.

[tianshilei1992]

I'm gonna accept it for now to make it in 14. We can always come back later if we see anything wrong.

clang/lib/Driver/ToolChains/Clang.cpp

void LinkerWrapper::ConstructJob(Compilation &C, const JobAction &JA,
   }
 
   // Construct the link job so we can wrap around it.
   Linker->ConstructJob(C, JA, Output, Inputs, Args, LinkingOutput);
   const auto &LinkCommand = C.getJobs().getJobs().back();
 
   CmdArgs.push_back("-linker-path");
   CmdArgs.push_back(LinkCommand->getExecutable());
+  CmdArgs.push_back("--");
   for (const char *LinkArg : LinkCommand->getArguments())
     CmdArgs.push_back(LinkArg);
 
   const char *Exec =
       Args.MakeArgString(getToolChain().GetProgramPath("clang-linker-wrapper"));
 
   // Replace the executable and arguments associated with the link job to the
   // wrapper.
   LinkCommand->replaceExecutable(Exec);
   LinkCommand->replaceArguments(CmdArgs);
 }

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

static cl::opt<std::string> OptLevel("opt-level",
                                      cl::init("O2"),
                                      cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<std::string>
     BitcodeLibrary("target-library",
                    cl::desc("Path for the target bitcode library"),
                    cl::cat(ClangLinkerWrapperCategory));
 
+static cl::opt<bool> EmbedBC(
+    "target-embed-bc", cl::ZeroOrMore,
+    cl::desc("Embed linked bitcode instead of an executable device image."),
+    cl::init(false), cl::cat(ClangLinkerWrapperCategory));
+
 // Do not parse linker options.
 static cl::list<std::string>
-    HostLinkerArgs(cl::Sink, cl::desc("<options to be passed to linker>..."));
+    HostLinkerArgs(cl::Positional,
+                   cl::desc("<options to be passed to linker>..."));
 
 /// Path of the current binary.
-static std::string LinkerExecutable;
+static const char *LinkerExecutable;
 
 /// Temporary files created by the linker wrapper.
 static SmallVector<std::string, 16> TempFiles;
 
 /// Codegen flags for LTO backend.
 static codegen::RegisterCodeGenFlags CodeGenFlags;
 
 /// Magic section string that marks the existence of offloading data. The
extractFromArchive(const Archive &Library,
   SmallString<128> TempFile;
   if (std::error_code EC =
           sys::fs::createTemporaryFile(Prefix + "-host", Extension, TempFile))
     return createFileError(TempFile, EC);
   TempFiles.push_back(static_cast<std::string>(TempFile));
 
   std::unique_ptr<MemoryBuffer> Buffer =
       MemoryBuffer::getMemBuffer(Library.getMemoryBufferRef(), false);
-  if (Error Err = writeArchive(TempFile, Members, true, Library.kind(),
-                                    true, Library.isThin(), std::move(Buffer)))
+  if (Error Err = writeArchive(TempFile, Members, true, Library.kind(), true,
+                               Library.isThin(), std::move(Buffer)))
     return std::move(Err);
 
   return static_cast<std::string>(TempFile);
 }
 
 /// Extracts embedded device offloading code from a memory \p Buffer to a list
 /// of \p DeviceFiles. If device code was extracted a new file with the embedded
 /// device code stripped from the buffer will be returned.
Expected<std::string> assemble(StringRef InputFile, Triple TheTriple,
   CmdArgs.push_back(InputFile);
 
   if (sys::ExecuteAndWait(*PtxasPath, CmdArgs))
     return createStringError(inconvertibleErrorCode(), "'ptxas' failed");
 
   return static_cast<std::string>(TempFile);
 }
 
-Expected<std::string> link(ArrayRef<StringRef> InputFiles,
+Expected<std::string> link(ArrayRef<std::string> InputFiles,
                            ArrayRef<std::string> LinkerArgs, Triple TheTriple,
                            StringRef Arch) {
   // NVPTX uses the nvlink binary to link device object files.
   ErrorOr<std::string> NvlinkPath = sys::findProgramByName("nvlink");
   if (!NvlinkPath)
     return createStringError(NvlinkPath.getError(),
                              "Unable to find 'nvlink' in path");
 
Expected<std::string> link(ArrayRef<std::string> InputFiles,
   CmdArgs.push_back(Arch);
 
   // Copy system library paths used by the host linker.
   for (StringRef Arg : LinkerArgs)
     if (Arg.startswith("-L"))
       CmdArgs.push_back(Arg);
 
   // Add extracted input files.
-  for (auto Input : InputFiles)
+  for (StringRef Input : InputFiles)
     CmdArgs.push_back(Input);
 
   if (sys::ExecuteAndWait(*NvlinkPath, CmdArgs))
     return createStringError(inconvertibleErrorCode(), "'nvlink' failed");
 
   return static_cast<std::string>(TempFile);
 }
 } // namespace nvptx
 
-Expected<std::string> linkDevice(ArrayRef<StringRef> InputFiles,
+Expected<std::string> linkDevice(ArrayRef<std::string> InputFiles,
                                  ArrayRef<std::string> LinkerArgs,
                                  Triple TheTriple, StringRef Arch) {
   switch (TheTriple.getArch()) {
   case Triple::nvptx:
   case Triple::nvptx64:
     return nvptx::link(InputFiles, LinkerArgs, TheTriple, Arch);
   case Triple::amdgcn:
     // TODO: AMDGCN linking support.
CodeGenOpt::Level getCGOptLevel(unsigned OptLevel) {
   case 2:
     return CodeGenOpt::Default;
   case 3:
     return CodeGenOpt::Aggressive;
   }
   llvm_unreachable("Invalid optimization level");
 }
 
-std::unique_ptr<lto::LTO> createLTO(const Triple &TheTriple, StringRef Arch,
-                                    bool WholeProgram) {
+template <typename ModuleHook = function_ref<bool(size_t, const Module &)>>
+std::unique_ptr<lto::LTO> createLTO(
+    const Triple &TheTriple, StringRef Arch, bool WholeProgram,
+    ModuleHook Hook = [](size_t, const Module &) { return true; }) {
   lto::Config Conf;
   lto::ThinBackend Backend;
   // TODO: Handle index-only thin-LTO
   Backend = lto::createInProcessThinBackend(
       llvm::heavyweight_hardware_concurrency(1));
 
   Conf.CPU = Arch.str();
   Conf.Options = codegen::InitTargetOptionsFromCodeGenFlags(TheTriple);
 
   Conf.MAttrs = getTargetFeatures(TheTriple);
   Conf.CGOptLevel = getCGOptLevel(OptLevel[1] - '0');
   Conf.OptLevel = OptLevel[1] - '0';
   Conf.DefaultTriple = TheTriple.getTriple();
   Conf.DiagHandler = diagnosticHandler;
 
   Conf.PTO.LoopVectorization = Conf.OptLevel > 1;
   Conf.PTO.SLPVectorization = Conf.OptLevel > 1;
 
-  // TODO: Handle outputting bitcode using a module hook.
+  Conf.PostInternalizeModuleHook = Hook;
   if (TheTriple.isNVPTX())
     Conf.CGFileType = CGFT_AssemblyFile;
   else
     Conf.CGFileType = CGFT_ObjectFile;
 
   // TODO: Handle remark files
   Conf.HasWholeProgramVisibility = WholeProgram;
 
   return std::make_unique<lto::LTO>(std::move(Conf), Backend);
 }
 
 // Returns true if \p S is valid as a C language identifier and will be given
 // `__start_` and `__stop_` symbols.
 bool isValidCIdentifier(StringRef S) {
   return !S.empty() && (isAlpha(S[0]) || S[0] == '_') &&
          std::all_of(S.begin() + 1, S.end(),
                      [](char C) { return C == '_' || isAlnum(C); });
 }
 
-Expected<Optional<std::string>> linkBitcodeFiles(ArrayRef<StringRef> InputFiles,
-                                                 const Triple &TheTriple,
-                                                 StringRef Arch) {
+Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
+                       const Triple &TheTriple, StringRef Arch) {
   SmallVector<std::unique_ptr<MemoryBuffer>, 4> SavedBuffers;
   SmallVector<std::unique_ptr<lto::InputFile>, 4> BitcodeFiles;
+  SmallVector<std::string, 4> NewInputFiles;
   StringMap<bool> UsedInRegularObj;
 
   // Search for bitcode files in the input and create an LTO input file. If it
   // is not a bitcode file, scan its symbol table for symbols we need to
   // save.
   for (StringRef File : InputFiles) {
     ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
         MemoryBuffer::getFileOrSTDIN(File);
     if (std::error_code EC = BufferOrErr.getError())
       return createFileError(File, EC);
 
     file_magic Type = identify_magic((*BufferOrErr)->getBuffer());
     if (Type != file_magic::bitcode) {
       Expected<std::unique_ptr<ObjectFile>> ObjFile =
           ObjectFile::createObjectFile(**BufferOrErr, Type);
       if (!ObjFile)
         return ObjFile.takeError();
 
+      NewInputFiles.push_back(File.str());
       for (auto &Sym : (*ObjFile)->symbols()) {
         Expected<StringRef> Name = Sym.getName();
         if (!Name)
           return Name.takeError();
 
         UsedInRegularObj[*Name] = true;
       }
     } else {
       Expected<std::unique_ptr<lto::InputFile>> InputFileOrErr =
           llvm::lto::InputFile::create(**BufferOrErr);
       if (!InputFileOrErr)
         return InputFileOrErr.takeError();
 
       BitcodeFiles.push_back(std::move(*InputFileOrErr));
       SavedBuffers.push_back(std::move(*BufferOrErr));
     }
   }
 
   if (BitcodeFiles.empty())
-    return None;
+    return Error::success();
+
+  auto HandleError = [&](std::error_code EC) {
+    logAllUnhandledErrors(errorCodeToError(EC),
+                          WithColor::error(errs(), LinkerExecutable));
+    exit(1);
+  };
+
+  // LTO Module hook to output bitcode without running the backend.
+  auto LinkOnly = [&](size_t Task, const Module &M) {
+    SmallString<128> TempFile;
+    if (std::error_code EC = sys::fs::createTemporaryFile(
+            "jit-" + TheTriple.getTriple(), "bc", TempFile))
+      HandleError(EC);
+    std::error_code EC;
+    raw_fd_ostream LinkedBitcode(TempFile, EC, sys::fs::OF_None);
+    if (EC)
+      HandleError(EC);
+    WriteBitcodeToFile(M, LinkedBitcode);
+    TempFiles.push_back(static_cast<std::string>(TempFile));
+    NewInputFiles.push_back(static_cast<std::string>(TempFile));
+    return false;
+  };
 
   // We have visibility of the whole program if every input is bitcode, all
   // inputs are statically linked so there should be no external references.
   bool WholeProgram = BitcodeFiles.size() == InputFiles.size();
-  StringMap<bool> PrevailingSymbols;
+  auto LTOBackend = (EmbedBC)
+                        ? createLTO(TheTriple, Arch, WholeProgram, LinkOnly)
+                        : createLTO(TheTriple, Arch, WholeProgram);
 
   // TODO: Run more tests to verify that this is correct.
   // Create the LTO instance with the necessary config and add the bitcode files
   // to it after resolving symbols. We make a few assumptions about symbol
   // resolution.
   // 1. The target is going to be a stand-alone executable file.
   // 2. We do not support relocatable object files.
   // 3. All inputs are relocatable object files extracted from host binaries, so
   //    there is no resolution to a dynamic library.
-  auto LTOBackend = createLTO(TheTriple, Arch, WholeProgram);
+  StringMap<bool> PrevailingSymbols;
   for (auto &BitcodeFile : BitcodeFiles) {
     const auto Symbols = BitcodeFile->symbols();
     SmallVector<lto::SymbolResolution, 16> Resolutions(Symbols.size());
     size_t Idx = 0;
     for (auto &Sym : Symbols) {
       lto::SymbolResolution &Res = Resolutions[Idx++];
 
       // We will use this as the prevailing symbol definition in LTO unless
Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
   size_t MaxTasks = LTOBackend->getMaxTasks();
   std::vector<SmallString<128>> Files(MaxTasks);
   auto AddStream = [&](size_t Task) -> std::unique_ptr<CachedFileStream> {
     int FD = -1;
     auto &TempFile = Files[Task];
     StringRef Extension = (TheTriple.isNVPTX()) ? "s" : "o";
     if (std::error_code EC = sys::fs::createTemporaryFile(
             "lto-" + TheTriple.getTriple(), Extension, FD, TempFile))
-      return nullptr;
+      HandleError(EC);
     TempFiles.push_back(static_cast<std::string>(TempFile));
     return std::make_unique<CachedFileStream>(
         std::make_unique<llvm::raw_fd_ostream>(FD, true));
   };
+
   if (Error Err = LTOBackend->run(AddStream))
     return std::move(Err);
 
+  // Is we are compiling for NVPTX we need to run the assembler first.
   for (auto &File : Files) {
-    if (!TheTriple.isNVPTX())
+    if (!TheTriple.isNVPTX() || EmbedBC)
       continue;
 
     auto FileOrErr = nvptx::assemble(File, TheTriple, Arch);
     if (!FileOrErr)
       return FileOrErr.takeError();
     File = *FileOrErr;
   }
 
-  return static_cast<std::string>(Files.front());
+  // Append the new inputs to the device linker input.
+  for (auto &File : Files)
+    NewInputFiles.push_back(static_cast<std::string>(File));
+  InputFiles = NewInputFiles;
+
+  return Error::success();
 }
 
 /// Runs the appropriate linking action on all the device files specified in \p
 /// DeviceFiles. The linked device images are returned in \p LinkedImages.
 Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
                       ArrayRef<std::string> LinkerArgs,
                       SmallVectorImpl<std::string> &LinkedImages) {
   // Get the list of inputs for a specific device.
-  StringMap<SmallVector<StringRef, 4>> LinkerInputMap;
+  StringMap<SmallVector<std::string, 4>> LinkerInputMap;
   for (auto &File : DeviceFiles)
     LinkerInputMap[StringRef(File)].push_back(File.Filename);
 
   // Try to link each device toolchain.
   for (auto &LinkerInput : LinkerInputMap) {
     auto TargetFeatures = LinkerInput.getKey().rsplit('-');
     Triple TheTriple(TargetFeatures.first);
     StringRef Arch(TargetFeatures.second);
 
-    // TODO: Run LTO or bitcode linking before the final link job.
-    auto ObjectOrErr =
-        linkBitcodeFiles(LinkerInput.getValue(), TheTriple, Arch);
-    if (!ObjectOrErr)
-      return ObjectOrErr.takeError();
-    if ((*ObjectOrErr).hasValue())
-      LinkerInput.getValue() = {**ObjectOrErr};
+    // Run LTO on any bitcode files and replace the input with the result.
+    if (Error Err = linkBitcodeFiles(LinkerInput.getValue(), TheTriple, Arch))
+      return std::move(Err);
+
+    // If we are embedding bitcode for JIT, skip the final device linking.
+    if (EmbedBC) {
+      assert(!LinkerInput.getValue().empty() && "No bitcode image to embed");
+      LinkedImages.push_back(LinkerInput.getValue().front());
+      continue;
+    }
 
     auto ImageOrErr =
         linkDevice(LinkerInput.getValue(), LinkerArgs, TheTriple, Arch);
     if (!ImageOrErr)
       return ImageOrErr.takeError();
 
     LinkedImages.push_back(*ImageOrErr);
   }