[OpenMP] Add basic support for properly handling static libraries

[jhuber6]

Summary

Currently we handle static libraries like any other object in the
linker wrapper. However, this does not preserve the sematnics that
dictate static libraries should be lazily loaded as the symbols are
needed. This allows us to ignore linking in architectures that are not
used by the main application being compiled. This patch adds the basic
support for detecting if a file came from a static library, and only
including it in the link job if it's used by other object files.

This patch only adds the basic support, to be more correct we should
check the symbols and only inclue the library if the link job contains
symbols that are needed. Ideally we could just put this on the linker
itself, but nvlink doesn't seem to support .a files.

[jdoerfert]

Typo(s) in the commit message, otherwise this seems reasonable. LG

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

1153

Nit: count?

[saiislam]

Ideally we could just put this on the linker itself, but nvlink doesn't seem to support .a files.

Yes, nvlink does not support archives. So we used a wrapper to extract cubin files from the archive and pass them to nvlink. Please see, Clang Nvlink Wrapper.

[jhuber6]

In D125092#3497033, @saiislam wrote:

Ideally we could just put this on the linker itself, but nvlink doesn't seem to support .a files.

Yes, nvlink does not support archives. So we used a wrapper to extract cubin files from the archive and pass them to nvlink. Please see, Clang Nvlink Wrapper.

Yeah, it seems we'll need to do a manual work-around later to handle the following case:

// foo.o
void foo() { return; }

//libfoo.a
void foo() { return; }

clang foo.o libfoo.a // no conflict, symbol foo isn't undefined in foo.o so libfoo.a is never loaded

nvlink foo.cubin libfoo.cubin // conflict, both are loaded and linked together

This is a little bit of an edge-case so I just went for the basic support here, since AMD just uses lld we could just pass in libfoo.a and it should handle it correctly for us. Add it to the list of Nvidia problems we need to work around externally.

clang/test/Driver/linker-wrapper.c

 
 // RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o \
 // RUN:   -fembed-offload-object=%S/Inputs/dummy-elf.o,openmp,nvptx64-nvida-cuda,sm_70 \
 // RUN:   -fembed-offload-object=%S/Inputs/dummy-elf.o,cuda,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=CUDA_OMP_LINK
 
 // CUDA_OMP_LINK: nvlink{{.*}}-m64 -o {{.*}}.out -arch sm_70 {{.*}}.o {{.*}}.o
+
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t-lib.o \
+// RUN:   -fembed-offload-object=%S/Inputs/dummy-elf.o,openmp,nvptx64-nvida-cuda,sm_70 \
+// RUN:   -fembed-offload-object=%S/Inputs/dummy-elf.o,openmp,nvptx64-nvida-cuda,sm_52
+// RUN: llvm-ar rcs %t.a %t-lib.o
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t-obj.o \
+// RUN:   -fembed-offload-object=%S/Inputs/dummy-elf.o,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.a %t-obj.o -o a.out 2>&1 | FileCheck %s --check-prefix=STATIC-LIBRARY
+
+// STATIC-LIBRARY: nvlink{{.*}} -arch sm_70
+// STATIC-LIBRARY-NOT: nvlink{{.*}} -arch sm_50

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

 
 /// The magic offset for the first object inside CUDA's fatbinary. This can be
 /// different but it should work for what is passed here.
 static constexpr unsigned FatbinaryOffset = 0x50;
 
 /// Information for a device offloading file extracted from the host.
 struct DeviceFile {
   DeviceFile(StringRef Kind, StringRef TheTriple, StringRef Arch,
-             StringRef Filename)
-      : Kind(Kind), TheTriple(TheTriple), Arch(Arch), Filename(Filename) {}
+             StringRef Filename, bool IsLibrary = false)
+      : Kind(Kind), TheTriple(TheTriple), Arch(Arch), Filename(Filename),
+        IsLibrary(IsLibrary) {}
 
   std::string Kind;
   std::string TheTriple;
   std::string Arch;
   std::string Filename;
+  bool IsLibrary;
 };
 
 namespace llvm {
 /// Helper that allows DeviceFile to be used as a key in a DenseMap. For now we
 /// assume device files with matching architectures and triples but different
 /// offloading kinds should be handlded together, this may not be true in the
 /// future.
 template <> struct DenseMapInfo<DeviceFile> {
template <> struct DenseMapInfo<DeviceFile> {
   }
 };
 } // namespace llvm
 
 namespace {
 
 Expected<Optional<std::string>>
 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
-                  SmallVectorImpl<DeviceFile> &DeviceFiles);
+                  SmallVectorImpl<DeviceFile> &DeviceFiles,
+                  bool IsLibrary = false);
 
 void printCommands(ArrayRef<StringRef> CmdArgs) {
   if (CmdArgs.empty())
     return;
 
   llvm::errs() << " \"" << CmdArgs.front() << "\" ";
   for (auto IC = std::next(CmdArgs.begin()), IE = CmdArgs.end(); IC != IE; ++IC)
     llvm::errs() << *IC << (std::next(IC) != IE ? " " : "\n");
GV = new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
                                 "llvm.compiler.used");
   GV->setSection("llvm.metadata");
 }
 
 /// Attempts to extract all the embedded device images contained inside the
 /// buffer \p Contents. The buffer is expected to contain a valid offloading
 /// binary format.
 Error extractOffloadFiles(StringRef Contents, StringRef Prefix,
-                          SmallVectorImpl<DeviceFile> &DeviceFiles) {
+                          SmallVectorImpl<DeviceFile> &DeviceFiles,
+                          bool IsLibrary = false) {
   uint64_t Offset = 0;
   // There could be multiple offloading binaries stored at this section.
   while (Offset < Contents.size()) {
     std::unique_ptr<MemoryBuffer> Buffer =
         MemoryBuffer::getMemBuffer(Contents.drop_front(Offset), "",
                                    /*RequiresNullTerminator*/ false);
     auto BinaryOrErr = OffloadBinary::create(*Buffer);
     if (!BinaryOrErr)
if (!OutputOrErr)
       return OutputOrErr.takeError();
     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
     std::copy(Binary.getImage().bytes_begin(), Binary.getImage().bytes_end(),
               Output->getBufferStart());
     if (Error E = Output->commit())
       return E;
 
     DeviceFiles.emplace_back(Kind, Binary.getTriple(), Binary.getArch(),
-                             TempFile);
+                             TempFile, IsLibrary);
 
     Offset += Binary.getSize();
   }
 
   return Error::success();
 }
 
 Expected<Optional<std::string>>
 extractFromBinary(const ObjectFile &Obj,
-                  SmallVectorImpl<DeviceFile> &DeviceFiles) {
+                  SmallVectorImpl<DeviceFile> &DeviceFiles,
+                  bool IsLibrary = false) {
   StringRef Extension = sys::path::extension(Obj.getFileName()).drop_front();
   StringRef Prefix = sys::path::stem(Obj.getFileName());
   SmallVector<StringRef, 4> ToBeStripped;
 
   // Extract offloading binaries from sections with the name `.llvm.offloading`.
   for (const SectionRef &Sec : Obj.sections()) {
     Expected<StringRef> Name = Sec.getName();
     if (!Name || !Name->equals(OFFLOAD_SECTION_MAGIC_STR))
       continue;
 
     Expected<StringRef> Contents = Sec.getContents();
     if (!Contents)
       return Contents.takeError();
 
-    if (Error Err = extractOffloadFiles(*Contents, Prefix, DeviceFiles))
+    if (Error Err =
+            extractOffloadFiles(*Contents, Prefix, DeviceFiles, IsLibrary))
       return std::move(Err);
 
     ToBeStripped.push_back(*Name);
   }
 
   if (ToBeStripped.empty() || !StripSections)
     return None;
 
extractFromBinary(const ObjectFile &Obj,
   if (Error Err = executeCommands(*StripPath, StripArgs))
     return std::move(Err);
 
   return static_cast<std::string>(TempFile);
 }
 
 Expected<Optional<std::string>>
 extractFromBitcode(std::unique_ptr<MemoryBuffer> Buffer,
-                   SmallVectorImpl<DeviceFile> &DeviceFiles) {
+                   SmallVectorImpl<DeviceFile> &DeviceFiles,
+                   bool IsLibrary = false) {
   LLVMContext Context;
   SMDiagnostic Err;
   std::unique_ptr<Module> M = getLazyIRModule(std::move(Buffer), Err, Context);
   if (!M)
     return createStringError(inconvertibleErrorCode(),
                              "Failed to create module");
 
   StringRef Extension = sys::path::extension(M->getName()).drop_front();
if (!GV.hasSection() || !GV.getSection().equals(OFFLOAD_SECTION_MAGIC_STR))
       continue;
 
     auto *CDS = dyn_cast<ConstantDataSequential>(GV.getInitializer());
     if (!CDS)
       continue;
 
     StringRef Contents = CDS->getAsString();
 
-    if (Error Err = extractOffloadFiles(Contents, Prefix, DeviceFiles))
+    if (Error Err =
+            extractOffloadFiles(Contents, Prefix, DeviceFiles, IsLibrary))
       return std::move(Err);
 
     ToBeDeleted.push_back(&GV);
   }
 
   if (ToBeDeleted.empty() || !StripSections)
     return None;
 
extractFromArchive(const Archive &Library,
   Error Err = Error::success();
   for (auto Child : Library.children(Err)) {
     auto ChildBufferRefOrErr = Child.getMemoryBufferRef();
     if (!ChildBufferRefOrErr)
       return ChildBufferRefOrErr.takeError();
     std::unique_ptr<MemoryBuffer> ChildBuffer =
         MemoryBuffer::getMemBuffer(*ChildBufferRefOrErr, false);
 
-    auto FileOrErr = extractFromBuffer(std::move(ChildBuffer), DeviceFiles);
+    auto FileOrErr = extractFromBuffer(std::move(ChildBuffer), DeviceFiles,
+                                       /*IsLibrary*/ true);
     if (!FileOrErr)
       return FileOrErr.takeError();
 
     // If we created a new stripped host file, use it to create a new archive
     // member, otherwise use the old member.
     if (!FileOrErr->hasValue()) {
       Expected<NewArchiveMember> NewMember =
           NewArchiveMember::getOldMember(Child, true);
extractFromArchive(const Archive &Library,
   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>>
 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
-                  SmallVectorImpl<DeviceFile> &DeviceFiles) {
+                  SmallVectorImpl<DeviceFile> &DeviceFiles, bool IsLibrary) {
   file_magic Type = identify_magic(Buffer->getBuffer());
   switch (Type) {
   case file_magic::bitcode:
-    return extractFromBitcode(std::move(Buffer), DeviceFiles);
+    return extractFromBitcode(std::move(Buffer), DeviceFiles, IsLibrary);
   case file_magic::elf_relocatable:
   case file_magic::macho_object:
   case file_magic::coff_object: {
     Expected<std::unique_ptr<ObjectFile>> ObjFile =
         ObjectFile::createObjectFile(*Buffer, Type);
     if (!ObjFile)
       return ObjFile.takeError();
-    return extractFromBinary(*ObjFile->get(), DeviceFiles);
+    return extractFromBinary(*ObjFile->get(), DeviceFiles, IsLibrary);
   }
   case file_magic::archive: {
     Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
         object::Archive::create(*Buffer);
     if (!LibFile)
       return LibFile.takeError();
     return extractFromArchive(*LibFile->get(), DeviceFiles);
   }

 }
 
 /// 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,
                       SmallVectorImpl<std::string> &LinkedImages) {
   // Get the list of inputs for a specific device.
   DenseMap<DeviceFile, SmallVector<std::string, 4>> LinkerInputMap;
-  for (auto &File : DeviceFiles)
+  SmallVector<DeviceFile, 4> LibraryFiles;
+  for (auto &File : DeviceFiles) {
+    if (File.IsLibrary)
+      LibraryFiles.push_back(File);
+    else
+      LinkerInputMap[File].push_back(File.Filename);
+  }
+
+  // Static libraries are loaded lazily as-needed, only add them if other files
+  // are present.
+  // TODO: We need to check the symbols as well, static libraries are only
+  //       loaded if they contain symbols that are currently undefined or common
+  //       in the symbol table.
+  for (auto &File : LibraryFiles)
+    if (LinkerInputMap.count(File))
       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.