[OpenMP] Link the bitcode library late for device LTO

[jhuber6]

Summary

This patch adds support for linking the OpenMP device bitcode library
late when doing LTO. This simply passes it in as an additional device
file when doing the final device linking phase with LTO. This has the
advantage that we don't link it multiple times, and the device
references do not get inlined and prevent us from doing needed OpenMP
optimizations when we have visiblity of the whole module.

Depends on D116975

[JonChesterfield]

Do we have a hook where we could link it at, uh, link time on nvtptx without LTO? Amdgpu had a llvm-link already there. Always bothered me a little that we link a copy per TU then hope the optimiser sorts it out nicely.

[jhuber6]

In D117048#3235763, @JonChesterfield wrote:

Do we have a hook where we could link it at, uh, link time on nvtptx without LTO? Amdgpu had a llvm-link already there. Always bothered me a little that we link a copy per TU then hope the optimiser sorts it out nicely.

I think it would be difficult, we could compile the library stand-alone and link it normally in the linker wrapper, but the performance would probably be bad. We could potentially only link everything once, but when we compile we don't really know how many files will need it right? It would definitely improve compile times if we could somehow do it though.

[JonChesterfield]

Oh, right. Nvptx is still lowering to machine code per-tu. We don't want the devicertl linking as machine code, so it has to go in per-tu. Or we could link nvptx as IR instead, and send that plus amdgpu down the same code path. Probably makes applications faster and compiles slower. Which sort of brings us to this patch with a different default.

[JonChesterfield]

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

289

Can/should probably always link it late on amdgpu

clang/lib/Driver/ToolChains/Clang.cpp

8176

This looks familiar - is there a function in commonargs that already does this?

[jhuber6]

clang/lib/Driver/ToolChains/Clang.cpp

8176

It's implemented individually when we select it in the CUDA and AMDGPUOpenMP toolchains, could probably try to unify it into a utility function to save a few lines.

[jhuber6]

Squash other uncommitted changes.

[tianshilei1992]

LGTM with two nits.

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

75

target-library is not the common name we call it. Maybe device-runtime-library?

967

Is this still needed now?

[jhuber6]

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

75

I think I called it bitcode-library later, since technically it's just any bitcode file that's linked in along LTO.

967

It still breaks, but I removed it later. My ordering for these patches is a bit of a mess, sorry.

clang/lib/Driver/ToolChains/AMDGPUOpenMP.cpp

void AMDGPUOpenMPToolChain::addClangTargetOptions(
 
   CC1Args.push_back("-target-cpu");
   CC1Args.push_back(DriverArgs.MakeArgStringRef(GPUArch));
   CC1Args.push_back("-fcuda-is-device");
 
   if (DriverArgs.hasArg(options::OPT_nogpulib))
     return;
 
+  // Link the bitcode library late if we're using device LTO.
+  if (getDriver().isUsingLTO(/* IsOffload */ true))
+    return;
+
   std::string BitcodeSuffix;
   if (DriverArgs.hasFlag(options::OPT_fopenmp_target_new_runtime,
                          options::OPT_fno_openmp_target_new_runtime, true))
     BitcodeSuffix = "new-amdgpu-" + GPUArch;
   else
     BitcodeSuffix = "amdgcn-" + GPUArch;
 
   addOpenMPDeviceRTL(getDriver(), DriverArgs, CC1Args, BitcodeSuffix,

clang/lib/Driver/ToolChains/Clang.cpp

for (auto TI = OpenMPTCRange.first, TE = OpenMPTCRange.second; TI != TE;
       const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
       ArgStringList FeatureArgs;
       TC->addClangTargetOptions(TCArgs, FeatureArgs, Action::OFK_OpenMP);
       auto FeatureIt = llvm::find(FeatureArgs, "-target-feature");
       CmdArgs.push_back(Args.MakeArgString(
           "-target-feature=" + TC->getTripleString() + "=" + *(FeatureIt + 1)));
     }
 
+    // Pass in the bitcode library to be linked during LTO.
+    for (auto TI = OpenMPTCRange.first, TE = OpenMPTCRange.second; TI != TE;
+         ++TI) {
+      const ToolChain *TC = TI->second;
+      const Driver &D = TC->getDriver();
+      const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
+      StringRef Arch = TCArgs.getLastArgValue(options::OPT_march_EQ);
+
+      std::string BitcodeSuffix;
+      if (TCArgs.hasFlag(options::OPT_fopenmp_target_new_runtime,
+                         options::OPT_fno_openmp_target_new_runtime, true))
+        BitcodeSuffix += "new-";
+      if (TC->getTriple().isNVPTX())
+        BitcodeSuffix += "nvptx-";
+      else if (TC->getTriple().isAMDGPU())
+        BitcodeSuffix += "amdgpu-";
+      BitcodeSuffix += Arch;
+
+      ArgStringList BitcodeLibrary;
+      addOpenMPDeviceRTL(D, TCArgs, BitcodeLibrary, BitcodeSuffix,
+                         TC->getTriple());
+
+      if (!BitcodeLibrary.empty())
+        CmdArgs.push_back(
+            Args.MakeArgString("-target-library=" + TC->getTripleString() +
+                               "-" + Arch + "=" + BitcodeLibrary.back()));
+    }
+
     // Pass in the optimization level to use for LTO.
     if (const Arg *A = Args.getLastArg(options::OPT_O_Group)) {
       StringRef OOpt;
       if (A->getOption().matches(options::OPT_O4) ||
           A->getOption().matches(options::OPT_Ofast))
         OOpt = "3";
       else if (A->getOption().matches(options::OPT_O)) {
         OOpt = A->getValue();

clang/lib/Driver/ToolChains/Cuda.cpp

#undef CASE_CUDA_VERSION
   if (DeviceOffloadingKind == Action::OFK_OpenMP) {
     if (CudaInstallationVersion < CudaVersion::CUDA_92) {
       getDriver().Diag(
           diag::err_drv_omp_offload_target_cuda_version_not_support)
           << CudaVersionToString(CudaInstallationVersion);
       return;
     }
 
+    // Link the bitcode library late if we're using device LTO.
+    if (getDriver().isUsingLTO(/* IsOffload */ true))
+      return;
+
     std::string BitcodeSuffix;
     if (DriverArgs.hasFlag(options::OPT_fopenmp_target_new_runtime,
                            options::OPT_fno_openmp_target_new_runtime, true))
       BitcodeSuffix = "new-nvptx-" + GpuArch.str();
     else
       BitcodeSuffix = "nvptx-" + GpuArch.str();
 
     addOpenMPDeviceRTL(getDriver(), DriverArgs, CC1Args, BitcodeSuffix,

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

static cl::opt<std::string> LinkerUserPath("linker-path",
                                            cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<std::string>
     TargetFeatures("target-feature", cl::desc("Target features for triple"),
                    cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<std::string> OptLevel("opt-level",
                                      cl::desc("Optimization level for LTO"),
-                                     cl::init("O0"),
-                                     cl::cat(ClangLinkerWrapperCategory));
+                                     cl::init("O2"),
+                                     cl::cat(ClangLinkerWrapperCategory));
+
+static cl::opt<std::string>
+    BitcodeLibrary("target-library",
+                   cl::desc("Path for the target bitcode library"),
 
 // Do not parse linker options.
 static cl::list<std::string>
     HostLinkerArgs(cl::Sink, cl::desc("<options to be passed to linker>..."));
 
 /// Path of the current binary.
 static std::string LinkerExecutable;
 
if (Expected<StringRef> Contents = Sec.getContents()) {
 
       Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
           FileOutputBuffer::create(TempFile, Sec.getSize());
       if (!OutputOrErr)
         return OutputOrErr.takeError();
       std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
       std::copy(Contents->begin(), Contents->end(), Output->getBufferStart());
       if (Error E = Output->commit())
-        return E;
+        return std::move(E);
 
       DeviceFiles.emplace_back(DeviceTriple, Arch, TempFile);
       ToBeStripped.push_back(*Name);
     }
   }
 
   if (ToBeStripped.empty())
     return None;
if (!sys::fs::exists(StripFile)) {
     auto Contents = Obj.getMemoryBufferRef().getBuffer();
     Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
         FileOutputBuffer::create(TempFile, Contents.size());
     if (!OutputOrErr)
       return OutputOrErr.takeError();
     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
     std::copy(Contents.begin(), Contents.end(), Output->getBufferStart());
     if (Error E = Output->commit())
-      return E;
+      return std::move(E);
     StripFile = TempFile;
   }
 
   // We will use llvm-strip to remove the now unneeded section containing the
   // offloading code.
   ErrorOr<std::string> StripPath = sys::findProgramByName("llvm-strip");
   if (!StripPath)
     return createStringError(StripPath.getError(),
for (GlobalVariable &GV : M->globals()) {
 
     Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
         FileOutputBuffer::create(TempFile, Contents.size());
     if (!OutputOrErr)
       return OutputOrErr.takeError();
     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
     std::copy(Contents.begin(), Contents.end(), Output->getBufferStart());
     if (Error E = Output->commit())
-      return E;
+      return std::move(E);
 
     DeviceFiles.emplace_back(DeviceTriple, Arch, TempFile);
     ToBeDeleted.push_back(&GV);
   }
 
   if (ToBeDeleted.empty())
     return None;
 
   // We need to materialize the lazy module before we make any changes.
   if (Error Err = M->materializeAll())
-    return Err;
+    return std::move(Err);
 
   // Remove the global from the module and write it to a new file.
   for (GlobalVariable *GV : ToBeDeleted) {
     removeFromCompilerUsed(*M, *GV);
     GV->eraseFromParent();
   }
 
   SmallString<128> TempFile;
if (!FileOrErr->hasValue()) {
 
       // We no longer need the stripped file, remove it.
       if (std::error_code EC = sys::fs::remove(**FileOrErr))
         return createFileError(**FileOrErr, EC);
     }
   }
 
   if (Err)
-    return Err;
+    return std::move(Err);
 
   if (!NewMembers)
     return None;
 
   // Create a new static library using the stripped host files.
   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 WriteErr = writeArchive(TempFile, Members, true, Library.kind(),
+  if (Error Err = writeArchive(TempFile, Members, true, Library.kind(),
                                     true, Library.isThin(), std::move(Buffer)))
-    return WriteErr;
+    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<Optional<std::string>>
for (auto &Sym : Symbols) {
       Res.LinkerRedefined = false;
 
       // Mark this symbol as the prevailing one.
       PrevailingSymbols[Sym.getName()] |= Res.Prevailing;
     }
 
     // Add the bitcode file with its resolved symbols to the LTO job.
     if (Error Err = LTOBackend->add(std::move(BitcodeFile), Resolutions))
-      return Err;
+      return std::move(Err);
   }
 
   // Run the LTO job to compile the bitcode.
   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;
     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 Err;
+    return std::move(Err);
 
   for (auto &File : Files) {
     if (!TheTriple.isNVPTX())
       continue;
 
     auto FileOrErr = nvptx::assemble(File, TheTriple, Arch);
     if (!FileOrErr)
       return FileOrErr.takeError();
if (sys::path::extension(Filename) == ".o" ||
       if (!NewFileOrErr)
         return reportError(NewFileOrErr.takeError());
 
       if (NewFileOrErr->hasValue())
         Arg = **NewFileOrErr;
     }
   }
 
+  // Add the device bitcode library to the device files if it was passed in.
+  if (!BitcodeLibrary.empty()) {
+    // FIXME: Hacky workaround to avoid a backend crash at O0.
+    if (OptLevel[1] - '0' == 0)
+      OptLevel[1] = '1';
+    auto DeviceAndPath = StringRef(BitcodeLibrary).split('=');
+    auto TripleAndArch = DeviceAndPath.first.rsplit('-');
+    DeviceFiles.emplace_back(TripleAndArch.first, TripleAndArch.second,
+                             DeviceAndPath.second);
+  }
+
   // Link the device images extracted from the linker input.
   SmallVector<std::string, 16> LinkedImages;
   if (Error Err = linkDeviceFiles(DeviceFiles, LinkerArgs, LinkedImages))
     return reportError(std::move(Err));
 
   // Wrap each linked device image into a linkable host binary and add it to the
   // link job's inputs.
   for (const auto &Image : LinkedImages) {