[OpenMP] Improve symbol resolution for OpenMP Offloading LTO

[jhuber6]

Summary

This patch improves the symbol resolution done for LTO with offloading
applications. The symbol resolution done here allows the LTO backend to
internalize more functions. The symbol resoltion done is a simplified
view that does not take into account various options like --wrap or
--dyanimic-list and always assumes we are creating a shared object.
The actual target may be an executable, but semantically it is used as a
shared object because certain objects need to be visible outside of the
executable when they are read by the OpenMP plugin.

Depends on D117246

[tianshilei1992]

LG with one nit.

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

763–768

[jhuber6]

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

763–768

Fixed.

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

 }
 
 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;
+  StringMap<bool> UsedInSharedLib;
 
   // 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();
 
+        // Record if we've seen these symbols in any object or shared libraries.
+        if ((*ObjFile)->isRelocatableObject()) {
           UsedInRegularObj[*Name] = true;
+        } else
+          UsedInSharedLib[*Name] = true;
       }
     } else {
       Expected<std::unique_ptr<lto::InputFile>> InputFileOrErr =
           llvm::lto::InputFile::create(**BufferOrErr);
       if (!InputFileOrErr)
         return InputFileOrErr.takeError();
 
+      // Save the input file and the buffer associated with its memory.
       BitcodeFiles.push_back(std::move(*InputFileOrErr));
       SavedBuffers.push_back(std::move(*BufferOrErr));
     }
   }
 
   if (BitcodeFiles.empty())
     return Error::success();
 
auto LinkOnly = [&](size_t Task, const Module &M) {
     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.
+  // We assume visibility of the whole program if every input file was bitcode.
   bool WholeProgram = BitcodeFiles.size() == InputFiles.size();
   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.
-  StringMap<bool> PrevailingSymbols;
+  // 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;
   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
-      // it is undefined in the module or another symbol has already been used.
-      Res.Prevailing = !Sym.isUndefined() && !PrevailingSymbols[Sym.getName()];
-
-      // We need LTO to preserve symbols referenced in other object files, or
-      // are needed by the rest of the toolchain.
+      // it is undefined or another definition has already been used.
+      Res.Prevailing =
+          !Sym.isUndefined() && PrevailingSymbols.insert(Sym.getName()).second;
+
+      // We need LTO to preseve the following global symbols:
+      // 1) Symbols used in regular objects.
+      // 2) Sections that will be given a __start/__stop symbol.
+      // 3) Prevailing symbols that are needed visibile to external libraries.
       Res.VisibleToRegularObj =
           UsedInRegularObj[Sym.getName()] ||
           isValidCIdentifier(Sym.getSectionName()) ||
-          (Res.Prevailing && Sym.getName().startswith("__omp"));
-
-      // We do not currently support shared libraries, so no symbols will be
-      // referenced externally by shared libraries.
-      Res.ExportDynamic = false;
-
-      // The result will currently always be an executable, so the only time the
-      // definition will not reside in this link unit is if it's undefined.
-      Res.FinalDefinitionInLinkageUnit = !Sym.isUndefined();
+          (Res.Prevailing &&
+           (Sym.getVisibility() != GlobalValue::HiddenVisibility &&
+            !Sym.canBeOmittedFromSymbolTable()));
+
+      // Identify symbols that must be exported dynamically and can be
+      // referenced by other files.
+      Res.ExportDynamic =
+          Sym.getVisibility() != GlobalValue::HiddenVisibility &&
+          (UsedInSharedLib[Sym.getName()] ||
+           !Sym.canBeOmittedFromSymbolTable());
+
+      // The final definition will reside in this linkage unit if the symbol is
+      // defined and local to the module. This only checks for bitcode files,
+      // full assertion will require complete symbol resolution.
+      Res.FinalDefinitionInLinkageUnit =
+          Sym.getVisibility() != GlobalValue::DefaultVisibility &&
+          (!Sym.isUndefined() && !Sym.isCommon());
 
       // We do not support linker redefined symbols (e.g. --wrap) for device
       // image linking, so the symbols will not be changed after LTO.
       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 std::move(Err);
+      return 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))
       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);
+    return Err;
 
   // Is we are compiling for NVPTX we need to run the assembler first.
   for (auto &File : Files) {
     if (!TheTriple.isNVPTX() || EmbedBC)
       continue;
 
     auto FileOrErr = nvptx::assemble(File, TheTriple, Arch);
     if (!FileOrErr)
Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
   // Try to link each device toolchain.
   for (auto &LinkerInput : LinkerInputMap) {
     auto TargetFeatures = LinkerInput.getKey().rsplit('-');
     Triple TheTriple(TargetFeatures.first);
     StringRef Arch(TargetFeatures.second);
 
     // 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);
+      return 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;
     }