[OpenMP] Use CUDA's non-RDC mode when LTO has whole program visibility

[jhuber6]

Summary

When we do LTO we consider ourselves to have whole program visibility if
every single input file we have contains LLVM bitcode. If we have whole
program visibliity then we can create a single image and utilize CUDA's
non-RDC mode by not passing -c to ptxas and ignoring the nvlink
job. This should be faster for some situations and also saves us the
time executing nvlink.

[tra]

LGTM with a minor test nit.

clang/test/Driver/linker-wrapper.c

41–42

// LTO-NOT: nvlink

[jhuber6]

Add test line

clang/test/Driver/linker-wrapper.c

 // HOST_LINK: ld.lld{{.*}}-a -b -c {{.*}}.o -o a.out
 
 // RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o \
 // RUN:   -fembed-offload-object=%S/Inputs/dummy-bc.bc,openmp,nvptx64-nvida-cuda,sm_70 \
 // RUN:   -fembed-offload-object=%S/Inputs/dummy-bc.bc,openmp,nvptx64-nvida-cuda,sm_70
 // RUN: clang-linker-wrapper --host-triple x86_64-unknown-linux-gnu --dry-run -linker-path \
 // RUN:   /usr/bin/ld -- %t.o -o a.out 2>&1 | FileCheck %s --check-prefix=LTO
 
-// LTO: ptxas{{.*}}-m64 -o {{.*}}.cubin -O2 --gpu-name sm_70 -c {{.*}}.s
-// LTO: nvlink{{.*}}-m64 -o {{.*}}.out -arch sm_70 {{.*}}.cubin
+// LTO: ptxas{{.*}}-m64 -o {{.*}}.cubin -O2 --gpu-name sm_70 {{.*}}.s
+// LTO-NOT: nvlink

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

extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
   default:
     return None;
   }
 }
 
 // TODO: Move these to a separate file.
 namespace nvptx {
 Expected<std::string> assemble(StringRef InputFile, Triple TheTriple,
-                               StringRef Arch) {
+                               StringRef Arch, bool RDC = true) {
   // NVPTX uses the ptxas binary to create device object files.
   Expected<std::string> PtxasPath = findProgram("ptxas", {CudaBinaryPath});
   if (!PtxasPath)
     return PtxasPath.takeError();
 
   // Create a new file to write the linked device image to.
   SmallString<128> TempFile;
   if (Error Err =
else if (DebugInfo == FullDebugInfo && OptLevel[1] == '0')
     CmdArgs.push_back("-g");
   for (auto &Arg : PtxasArgs)
     CmdArgs.push_back(Arg);
   CmdArgs.push_back("-o");
   CmdArgs.push_back(TempFile);
   CmdArgs.push_back(Opt);
   CmdArgs.push_back("--gpu-name");
   CmdArgs.push_back(Arch);
+  if (RDC)
     CmdArgs.push_back("-c");
 
   CmdArgs.push_back(InputFile);
 
   if (Verbose)
     printCommands(CmdArgs);
 
   if (Error Err = executeCommands(*PtxasPath, CmdArgs))
     return std::move(Err);

 // `__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); });
 }
 
 Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
-                       const Triple &TheTriple, StringRef Arch) {
+                       const Triple &TheTriple, StringRef Arch,
+                       bool &WholeProgram) {
   SmallVector<std::unique_ptr<MemoryBuffer>, 4> SavedBuffers;
   SmallVector<std::unique_ptr<lto::InputFile>, 4> BitcodeFiles;
   SmallVector<std::string, 4> NewInputFiles;
   DenseSet<StringRef> UsedInRegularObj;
   DenseSet<StringRef> UsedInSharedLib;
   BumpPtrAllocator Alloc;
   StringSaver Saver(Alloc);
 
auto OutputBitcode = [&](size_t Task, const Module &M) {
     if (EC)
       HandleError(errorCodeToError(EC));
     WriteBitcodeToFile(M, LinkedBitcode);
     NewInputFiles.push_back(static_cast<std::string>(TempFile));
     return false;
   };
 
   // We assume visibility of the whole program if every input file was bitcode.
-  bool WholeProgram = BitcodeFiles.size() == InputFiles.size();
+  WholeProgram = BitcodeFiles.size() == InputFiles.size();
   auto LTOBackend =
       (EmbedBitcode) ? createLTO(TheTriple, Arch, WholeProgram, OutputBitcode)
                      : createLTO(TheTriple, Arch, WholeProgram);
 
   // 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.
   DenseSet<StringRef> PrevailingSymbols;
Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
   };
 
   if (Error Err = LTOBackend->run(AddStream))
     return Err;
 
   // Is we are compiling for NVPTX we need to run the assembler first.
   if (TheTriple.isNVPTX() && !EmbedBitcode) {
     for (auto &File : Files) {
-      auto FileOrErr = nvptx::assemble(File, TheTriple, Arch);
+      auto FileOrErr = nvptx::assemble(File, TheTriple, Arch, !WholeProgram);
       if (!FileOrErr)
         return FileOrErr.takeError();
       File = *FileOrErr;
     }
   }
 
   // Append the new inputs to the device linker input.
   for (auto &File : Files)
Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
   DenseMap<DeviceFile, SmallVector<std::string, 4>> LinkerInputMap;
   for (auto &File : DeviceFiles)
     LinkerInputMap[File].push_back(File.Filename);
 
   // Try to link each device toolchain.
   for (auto &LinkerInput : LinkerInputMap) {
     DeviceFile &File = LinkerInput.getFirst();
     Triple TheTriple = Triple(File.TheTriple);
+    bool WholeProgram = false;
 
     // Run LTO on any bitcode files and replace the input with the result.
-    if (Error Err =
-            linkBitcodeFiles(LinkerInput.getSecond(), TheTriple, File.Arch))
+    if (Error Err = linkBitcodeFiles(LinkerInput.getSecond(), TheTriple,
+                                     File.Arch, WholeProgram))
       return Err;
 
     // If we are embedding bitcode for JIT, skip the final device linking.
     if (EmbedBitcode) {
       assert(!LinkerInput.getSecond().empty() && "No bitcode image to embed");
       LinkedImages.push_back(LinkerInput.getSecond().front());
       continue;
     }
 
+    // If we performed LTO on NVPTX and had whole program visibility, we can use
+    // CUDA in non-RDC mode.
+    if (WholeProgram && TheTriple.isNVPTX()) {
+      assert(!LinkerInput.getSecond().empty() && "No non-RDC image to embed");
+      LinkedImages.push_back(LinkerInput.getSecond().front());
+      continue;
+    }
+
     auto ImageOrErr = linkDevice(LinkerInput.getSecond(), TheTriple, File.Arch);
     if (!ImageOrErr)
       return ImageOrErr.takeError();
 
     LinkedImages.push_back(*ImageOrErr);
   }
   return Error::success();
 }