[LinkerWrapper] Fix resolution of weak symbols during LTO

[jhuber6]

Summary:
Weak symbols are supposed to have the semantics that they can be
overriden by a strong (i.e. global) definition. This wasn't being
respected by the LTO pass because we simply used the first definition
that was available. This patch fixes that logic by doing a first pass
over the symbols to check for strong resolutions that could override a
weak one.

A lot of fake linker logic is ending up in the linker wrapper. If there
were an option to handle this in lld it would be a lot cleaner, but
unfortunately supporting NVPTX is a big restriction as their binaries
require the nvlink tool.

[llvmbot]

@llvm/pr-subscribers-clang

Changes

Summary:
Weak symbols are supposed to have the semantics that they can be
overriden by a strong (i.e. global) definition. This wasn't being
respected by the LTO pass because we simply used the first definition
that was available. This patch fixes that logic by doing a first pass
over the symbols to check for strong resolutions that could override a
weak one.

A lot of fake linker logic is ending up in the linker wrapper. If there
were an option to handle this in lld it would be a lot cleaner, but
unfortunately supporting NVPTX is a big restriction as their binaries
require the nvlink tool.

---

Full diff: https://github.com/llvm/llvm-project/pull/68215.diff

2 Files Affected:

• (modified) clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp (+14)
• (added) openmp/libomptarget/test/offloading/weak.c (+32)

diff --git a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
index 632e37e3cac8fec..f95b0f8cb317c75 100644
--- a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -595,6 +595,7 @@ Error linkBitcodeFiles(SmallVectorImpl<OffloadFile> &InputFiles,
   StringRef Arch = Args.getLastArgValue(OPT_arch_EQ);
 
   SmallVector<OffloadFile, 4> BitcodeInputFiles;
+  DenseSet<StringRef> StrongResolutions;
   DenseSet<StringRef> UsedInRegularObj;
   DenseSet<StringRef> UsedInSharedLib;
   BumpPtrAllocator Alloc;
@@ -608,6 +609,18 @@ Error linkBitcodeFiles(SmallVectorImpl<OffloadFile> &InputFiles,
     file_magic Type = identify_magic(Buffer.getBuffer());
     switch (Type) {
     case file_magic::bitcode: {
+      Expected<IRSymtabFile> IRSymtabOrErr = readIRSymtab(Buffer);
+      if (!IRSymtabOrErr)
+        return IRSymtabOrErr.takeError();
+
+      // Check for any strong resolutions we need to preserve.
+      for (unsigned I = 0; I != IRSymtabOrErr->Mods.size(); ++I) {
+        for (const auto &Sym : IRSymtabOrErr->TheReader.module_symbols(I)) {
+          if (!Sym.isFormatSpecific() && Sym.isGlobal() && !Sym.isWeak() &&
+              !Sym.isUndefined())
+            StrongResolutions.insert(Saver.save(Sym.Name));
+        }
+      }
       BitcodeInputFiles.emplace_back(std::move(File));
       continue;
     }
@@ -696,6 +709,7 @@ Error linkBitcodeFiles(SmallVectorImpl<OffloadFile> &InputFiles,
       // it is undefined or another definition has already been used.
       Res.Prevailing =
           !Sym.isUndefined() &&
+          !(Sym.isWeak() && StrongResolutions.contains(Sym.getName())) &&
           PrevailingSymbols.insert(Saver.save(Sym.getName())).second;
 
       // We need LTO to preseve the following global symbols:
diff --git a/openmp/libomptarget/test/offloading/weak.c b/openmp/libomptarget/test/offloading/weak.c
new file mode 100644
index 000000000000000..b4f264cee7f2d70
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/weak.c
@@ -0,0 +1,32 @@
+// RUN: %libomptarget-compile-generic -DA -c -o %t-a.o
+// RUN: %libomptarget-compile-generic -DB -c -o %t-b.o
+// RUN: %libomptarget-compile-generic %t-a.o %t-b.o && %libomptarget-run-generic | %fcheck-generic
+
+#if defined(A)
+__attribute__((weak)) int x = 999;
+#pragma omp declare target to(x)
+#elif defined(B)
+int x = 42;
+#pragma omp declare target to(x)
+__attribute__((weak)) int y = 42;
+#pragma omp declare target to(y)
+#else
+
+#include <stdio.h>
+
+extern int x;
+#pragma omp declare target to(x)
+extern int y;
+#pragma omp declare target to(y)
+
+int main() {
+  x = 0;
+
+#pragma omp target update from(x)
+#pragma omp target update from(y)
+
+  // CHECK: PASS
+  if (x == 42 && y == 42)
+    printf("PASS\n");
+}
+#endif

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[JonChesterfield]

LinkerWrapper turning into a linker is kind of inevitable and not a very happy thing.

One option would be to lean on lld for amdgpu and split out the nvptx stuff in the hopes that we eventually have an alternative to nvlink, but it seems moderately unlikely to come to pass and driving both amdgpu and nvptx down the same code path has merits.