[Clang] Initial support for linking offloading code in tool

[jhuber6]

Summary

This patch adds the initial support for linking NVPTX offloading code
using the clang-linker-wrapper tool. This uses the extracted device
files and runs nvlink on them. Currently this is then passed to the
existing toolchain for creating linkable OpenMP offloading programs
using clang-offload-wrapper and compiling it manually using llc.
More work is required to support LTO, Bitcode linking, AMDGPU, and x86
offloading.

Depends on D116545

[jhuber6]

Small changes.

[jhuber6]

Need to tell llc to make position independent code so we can link it properly.

[jdoerfert]

Same as the other commit, where is this tested?

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

414

Unsure why we look at the linker bin dir first, PATH seems to be the right choice here.

428

If this is not addressed in a follow up, it seems trivial to add it.

460

keep exit success, makes it obvious this worked.

486

.str() above to get a key is somewhat ok, but then parsing the key again is weird.
Can we make it a densemap and provide a densemapinfo for struct DeviceFile. It can just inherit from the DenseMapInfo<std::string>. Might even work w/o explicitly calling the "super" functions with the .str() version if we have a std::string operator. Might then even work w/o DenseMapInfo by telling the DenseMap to use the std::string specialization of DenseMapInfo in the first place.

614

Try to replace auto where the type is short and helpful. E.g., here it would help to know what it is.

630

do we really want to return here?

[jhuber6]

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

414

Copied it from somewhere, should change it. Right now it's just a weird way to look through /bin/

428

It is addressed later.

460

Will do.

486

Probably a smarter idea, I needed a way to group all files with the same device file type. I'll look into it.

630

Might be a good idea to keep trying to remove temp files even if one fails, will address.

[jhuber6]

Maxing suggested changes.

[jdoerfert]

LG, as part of the patch set, and given the runtime test for coverage.

Some nits are left for later though.

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

 #include "llvm/Object/Archive.h"
 #include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Host.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
static cl::opt<bool> StripSections(
     cl::init(true), cl::cat(ClangLinkerWrapperCategory));
 
 static cl::opt<std::string> LinkerUserPath("linker-path",
                                            cl::desc("Path of linker binary"),
                                            cl::cat(ClangLinkerWrapperCategory));
 
 // Do not parse linker options.
 static cl::list<std::string>
-    LinkerArgs(cl::Sink, cl::desc("<options to be passed to linker>..."));
+    HostLinkerArgs(cl::Sink, cl::desc("<options to be passed to linker>..."));
 
 /// Path of the current binary.
 static std::string LinkerExecutable;
 
+static SmallVector<std::string, 16> TempFiles;
 /// Magic section string that marks the existence of offloading data. The
 /// section string will be formatted as `.llvm.offloading.<triple>.<arch>`.
-#define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading"
+#define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading."
 
+/// Information for a device offloading file extracted from the host.
 struct DeviceFile {
   DeviceFile(StringRef TheTriple, StringRef Arch, StringRef Filename)
       : TheTriple(TheTriple), Arch(Arch), Filename(Filename) {}
 
-  const Triple TheTriple;
+  const std::string TheTriple;
   const std::string Arch;
   const std::string Filename;
+
+  operator std::string() const { return TheTriple + "-" + Arch; }
 };
 
 namespace {
 
 Expected<Optional<std::string>>
 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
                   SmallVectorImpl<DeviceFile> &DeviceFiles);
 
+static StringRef getDeviceFileExtension(StringRef DeviceTriple,
+                                        bool IsBitcode = false) {
+  Triple TheTriple(DeviceTriple);
+  if (TheTriple.isAMDGPU() || IsBitcode)
+    return "bc";
+  if (TheTriple.isNVPTX())
+    return "cubin";
+  return "o";
+}
+
 Error runLinker(std::string &LinkerPath, SmallVectorImpl<std::string> &Args) {
   std::vector<StringRef> LinkerArgs;
   LinkerArgs.push_back(LinkerPath);
   for (auto &Arg : Args)
     LinkerArgs.push_back(Arg);
 
   if (sys::ExecuteAndWait(LinkerPath, LinkerArgs))
     return createStringError(inconvertibleErrorCode(), "'linker' failed");
for (const SectionRef &Sec : Obj.sections()) {
 
     SmallVector<StringRef, 4> SectionFields;
     Name->split(SectionFields, '.');
     StringRef DeviceTriple = SectionFields[3];
     StringRef Arch = SectionFields[4];
 
     if (Expected<StringRef> Contents = Sec.getContents()) {
       SmallString<128> TempFile;
+      StringRef DeviceExtension = getDeviceFileExtension(
+          DeviceTriple, identify_magic(*Contents) == file_magic::bitcode);
       if (std::error_code EC = sys::fs::createTemporaryFile(
-              Prefix + "-device-" + DeviceTriple, Extension, TempFile))
+              Prefix + "-device-" + DeviceTriple, DeviceExtension, TempFile))
         return createFileError(TempFile, EC);
+      TempFiles.push_back(static_cast<std::string>(TempFile));
 
       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;
 
       DeviceFiles.emplace_back(DeviceTriple, Arch, TempFile);
       ToBeStripped.push_back(*Name);
     }
   }
 
   if (ToBeStripped.empty())
     return None;
 
   // We will use llvm-strip to remove the now unneeded section containing the
   // offloading code.
-  ErrorOr<std::string> StripPath = sys::findProgramByName(
-      "llvm-strip", sys::path::parent_path(LinkerExecutable));
-  if (!StripPath)
-    StripPath = sys::findProgramByName("llvm-strip");
+  ErrorOr<std::string> StripPath = sys::findProgramByName("llvm-strip");
   if (!StripPath)
     return createStringError(StripPath.getError(),
                              "Unable to find 'llvm-strip' in path");
 
   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));
 
   SmallVector<StringRef, 8> StripArgs;
   StripArgs.push_back(*StripPath);
   StripArgs.push_back("--no-strip-all");
   StripArgs.push_back(Obj.getFileName());
   for (auto &Section : ToBeStripped) {
     StripArgs.push_back("--remove-section");
     StripArgs.push_back(Section);
for (GlobalVariable &GV : M->globals()) {
 
     SmallVector<StringRef, 4> SectionFields;
     GV.getSection().split(SectionFields, '.');
     StringRef DeviceTriple = SectionFields[3];
     StringRef Arch = SectionFields[4];
 
     StringRef Contents = CDS->getAsString();
     SmallString<128> TempFile;
+    StringRef DeviceExtension = getDeviceFileExtension(
+        DeviceTriple, identify_magic(Contents) == file_magic::bitcode);
     if (std::error_code EC = sys::fs::createTemporaryFile(
-            Prefix + "-device-" + DeviceTriple, Extension, TempFile))
+            Prefix + "-device-" + DeviceTriple, DeviceExtension, TempFile))
       return createFileError(TempFile, EC);
+    TempFiles.push_back(static_cast<std::string>(TempFile));
 
     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())
for (GlobalVariable *GV : ToBeDeleted) {
     removeFromCompilerUsed(*M, *GV);
     GV->eraseFromParent();
   }
 
   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::error_code EC;
   raw_fd_ostream HostOutput(TempFile, EC, sys::fs::OF_None);
   if (EC)
     return createFileError(TempFile, EC);
   WriteBitcodeToFile(*M, HostOutput);
   return static_cast<std::string>(TempFile);
 }
 
extractFromArchive(const Archive &Library,
   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(),
                                     true, Library.isThin(), std::move(Buffer)))
     return WriteErr;
 
   return static_cast<std::string>(TempFile);
Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
         object::Archive::create(*Buffer);
     if (!LibFile)
       return LibFile.takeError();
     return extractFromArchive(*LibFile->get(), DeviceFiles);
   }
   default:
     return errorCodeToError(object_error::invalid_file_type);
   }
+}
 
-  return None;
+// TODO: Move these to a separate file.
+namespace nvptx {
+Expected<std::string> link(ArrayRef<StringRef> 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");
+
+  // Create a new file to write the linked device image to.
+  SmallString<128> TempFile;
+  if (std::error_code EC = sys::fs::createTemporaryFile(
+          TheTriple.getArchName() + "-" + Arch + "-image", "out", TempFile))
+    return createFileError(TempFile, EC);
+  TempFiles.push_back(static_cast<std::string>(TempFile));
+
+  // TODO: Pass in arguments like `-g` and `-v` from the driver.
+  SmallVector<StringRef, 16> CmdArgs;
+  CmdArgs.push_back(*NvlinkPath);
+  CmdArgs.push_back(TheTriple.isArch64Bit() ? "-m64" : "-m32");
+  CmdArgs.push_back("-o");
+  CmdArgs.push_back(TempFile);
+  CmdArgs.push_back("-arch");
+  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)
+    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,
+                                 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.
+  case Triple::x86:
+  case Triple::x86_64:
+    // TODO: x86 linking support.
+  default:
+    return createStringError(inconvertibleErrorCode(),
+                             TheTriple.getArchName() +
+                                 " linking is not supported");
+  }
+}
+
+/// 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;
+  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 ImageOrErr =
+        linkDevice(LinkerInput.getValue(), LinkerArgs, TheTriple, Arch);
+    if (!ImageOrErr)
+      return ImageOrErr.takeError();
+
+    LinkedImages.push_back(*ImageOrErr);
+  }
+  return Error::success();
+}
+
+/// Creates an object file containing the device image stored in the filename \p
+/// ImageFile that can be linked with the host.
+Expected<std::string> wrapDeviceImage(StringRef ImageFile) {
+  // TODO: Call these utilities as a library intead of executing them here.
+  ErrorOr<std::string> WrapperPath =
+      sys::findProgramByName("clang-offload-wrapper");
+  if (!WrapperPath)
+    return createStringError(WrapperPath.getError(),
+                             "Unable to find 'clang-offload-wrapper' in path");
+
+  // Create a new file to write the wrapped bitcode file to.
+  SmallString<128> BitcodeFile;
+  if (std::error_code EC =
+          sys::fs::createTemporaryFile("offload", "bc", BitcodeFile))
+    return createFileError(BitcodeFile, EC);
+  TempFiles.push_back(static_cast<std::string>(BitcodeFile));
+
+  // TODO: Optionally pass the host triple in somewhere.
+  Triple HostTriple(sys::getDefaultTargetTriple());
+  SmallVector<StringRef, 4> WrapperArgs;
+  WrapperArgs.push_back(*WrapperPath);
+  WrapperArgs.push_back("-target");
+  WrapperArgs.push_back(HostTriple.getTriple());
+  WrapperArgs.push_back("-o");
+  WrapperArgs.push_back(BitcodeFile);
+  WrapperArgs.push_back(ImageFile);
+
+  if (sys::ExecuteAndWait(*WrapperPath, WrapperArgs))
+    return createStringError(inconvertibleErrorCode(),
+                             "'clang-offload-wrapper' failed");
+
+  ErrorOr<std::string> CompilerPath = sys::findProgramByName("llc");
+  if (!WrapperPath)
+    return createStringError(WrapperPath.getError(),
+                             "Unable to find 'llc' in path");
+
+  // Create a new file to write the wrapped bitcode file to.
+  SmallString<128> ObjectFile;
+  if (std::error_code EC =
+          sys::fs::createTemporaryFile("offload", "o", ObjectFile))
+    return createFileError(BitcodeFile, EC);
+  TempFiles.push_back(static_cast<std::string>(ObjectFile));
+
+  SmallVector<StringRef, 4> CompilerArgs;
+  CompilerArgs.push_back(*CompilerPath);
+  CompilerArgs.push_back("--filetype=obj");
+  CompilerArgs.push_back("--relocation-model=pic");
+  CompilerArgs.push_back("-o");
+  CompilerArgs.push_back(ObjectFile);
+  CompilerArgs.push_back(BitcodeFile);
+
+  if (sys::ExecuteAndWait(*CompilerPath, CompilerArgs))
+    return createStringError(inconvertibleErrorCode(), "'llc' failed");
+
+  return static_cast<std::string>(ObjectFile);
 }
 
 } // namespace
 
 int main(int argc, const char **argv) {
   InitLLVM X(argc, argv);
 
+  LinkerExecutable = argv[0];
   sys::PrintStackTraceOnErrorSignal(argv[0]);
   cl::SetVersionPrinter(PrintVersion);
   cl::HideUnrelatedOptions(ClangLinkerWrapperCategory);
   cl::ParseCommandLineOptions(
       argc, argv,
       "A wrapper utility over the host linker. It scans the input files for\n"
       "sections that require additional processing prior to linking. The tool\n"
       "will then transparently pass all arguments and input to the specified\n"
       "host linker to create the final binary.\n");
 
   if (Help) {
     cl::PrintHelpMessage();
     return EXIT_SUCCESS;
   }
-  LinkerExecutable = argv[0];
-
-  SmallVector<std::string, 4> TempFiles;
-  SmallVector<DeviceFile, 4> DeviceFiles;
 
   auto reportError = [argv](Error E) {
     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
-    exit(EXIT_FAILURE);
+    return EXIT_FAILURE;
   };
 
+  SmallVector<std::string, 16> LinkerArgs;
+  for (const std::string &Arg : HostLinkerArgs)
+    LinkerArgs.push_back(Arg);
+
   // Try to extract device code from the linker input and replace the linker
   // input with a new file that has the device section stripped.
+  SmallVector<DeviceFile, 4> DeviceFiles;
   for (std::string &Arg : LinkerArgs) {
     if (sys::path::extension(Arg) == ".o" ||
         sys::path::extension(Arg) == ".a") {
       ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
           MemoryBuffer::getFileOrSTDIN(Arg);
       if (std::error_code EC = BufferOrErr.getError())
-        reportError(createFileError(Arg, EC));
+        return reportError(createFileError(Arg, EC));
 
       auto NewFileOrErr =
           extractFromBuffer(std::move(*BufferOrErr), DeviceFiles);
 
       if (!NewFileOrErr)
-        reportError(NewFileOrErr.takeError());
+        return reportError(NewFileOrErr.takeError());
 
-      if (NewFileOrErr->hasValue()) {
-        TempFiles.push_back(**NewFileOrErr);
+      if (NewFileOrErr->hasValue())
         Arg = **NewFileOrErr;
-      }
-    }
-  }
 
-  // Add the newly extracted device files to the temporary list.
-  for (const auto &DeviceFile : DeviceFiles)
-    TempFiles.push_back(DeviceFile.Filename);
-
-  // TODO: Perform appropriate device linking action.
-  // TODO: Wrap device image in a host binary and pass it to the linker.
-  WithColor::warning(errs(), argv[0]) << "Offload linking not yet supported.\n";
-
-  SmallVector<std::string, 16> LinkerArgv;
-  for (const std::string &Arg : LinkerArgs)
-    LinkerArgv.push_back(Arg);
+  // 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) {
+    auto FileOrErr = wrapDeviceImage(Image);
+    if (!FileOrErr)
+      return reportError(FileOrErr.takeError());
+
+    LinkerArgs.push_back(*FileOrErr);
+  }
 
   // Run the host linking job.
-  if (Error Err = runLinker(LinkerUserPath, LinkerArgv))
-    reportError(std::move(Err));
+  if (Error Err = runLinker(LinkerUserPath, LinkerArgs))
+    return reportError(std::move(Err));
 
-  for (const auto &TempFile : TempFiles) {
+  // Remove the temporary files created.
+  for (const auto &TempFile : TempFiles)
     if (std::error_code EC = sys::fs::remove(TempFile))
       reportError(createFileError(TempFile, EC));
-  }
 
   return EXIT_SUCCESS;
 }