[Openmp-commits] [openmp] [LinkerWrapper] Fix resolution of weak symbols during LTO (PR #68215)

[Joseph Huber via Openmp-commits]

https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/68215

>From f2e29b6136d245e0712bfb87de1f4253d12dd6f5 Mon Sep 17 00:00:00 2001
From: Joseph Huber <jhuber6 at vols.utk.edu>
Date: Wed, 4 Oct 2023 07:34:01 -0500
Subject: [PATCH] [LinkerWrapper] Fix resolution of weak symbols during LTO

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.
---
 .../ClangLinkerWrapper.cpp                    | 14 ++++++++
 openmp/libomptarget/test/offloading/weak.c    | 33 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 openmp/libomptarget/test/offloading/weak.c

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..ca81db958356b2e
--- /dev/null
+++ b/openmp/libomptarget/test/offloading/weak.c
@@ -0,0 +1,33 @@
+// 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 && \
+// RUN:   %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