[NvlinkWrapper] Add support for `--undefined` (#113934)

Summary:
This flag is pretty canonical in ELF linkers, it allows us to force the
link job to extract a library if it defines a specific symbol. This is
mostly useful for letting us forcibly extract things that don't fit the
normal model (i.e. kernels) from static libraries.

3 files changed, 18 insertions, 5 deletions

diff --git a/clang/test/Driver/nvlink-wrapper.c b/clang/test/Driver/nvlink-wrapper.c
index 2b0993c..79f4a66 100644
--- a/clang/test/Driver/nvlink-wrapper.c
+++ b/clang/test/Driver/nvlink-wrapper.c
@@ -21,12 +21,13 @@ int bar() {
 }
 #else
 extern int y;
-int __attribute__((visibility("hidden"))) x = 999;
+extern int x;
 int baz() { return y + x; }
 #endif
 
 // Create various inputs to test basic linking and LTO capabilities. Creating a
 // CUDA binary requires access to the `ptxas` executable, so we just use x64.
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -o %t.o
 // RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -DX -o %t-x.o
 // RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -DY -o %t-y.o
 // RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -DZ -o %t-z.o
@@ -36,6 +37,7 @@ int baz() { return y + x; }
 // RUN: llvm-ar rcs %t-y.a %t-y.o
 // RUN: llvm-ar rcs %t-z.a %t-z.o
 // RUN: llvm-ar rcs %t-w.a %t-w.o
+// RUN: llvm-ar rcs %t-u.a %t-u.o
 
 //
 // Check that we forward any unrecognized argument to 'nvlink'.
@@ -49,11 +51,16 @@ int baz() { return y + x; }
 // `libx.a` and `liby.a` because extern weak symbols do not extract and `libz.a`
 // is not used at all.
 //
-// RUN: clang-nvlink-wrapper --dry-run %t-x.a %t-u.o %t-y.a %t-z.a %t-w.a \
+// RUN: clang-nvlink-wrapper --dry-run %t-x.a %t-u.a %t-y.a %t-z.a %t-w.a %t.o \
 // RUN:   -arch sm_52 -o a.out 2>&1 | FileCheck %s --check-prefix=LINK
 // LINK: nvlink{{.*}} -arch sm_52 -o a.out [[INPUT:.+]].cubin {{.*}}-x-{{.*}}.cubin{{.*}}-y-{{.*}}.cubin
 
-// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -o %t.o
+//
+// Same as above but we use '--undefined' to forcibly extract 'libz.a'
+//
+// RUN: clang-nvlink-wrapper --dry-run %t-x.a %t-u.a %t-y.a %t-z.a %t-w.a %t.o \
+// RUN:   -u z -arch sm_52 -o a.out 2>&1 | FileCheck %s --check-prefix=LINK
+// UNDEFINED: nvlink{{.*}} -arch sm_52 -o a.out [[INPUT:.+]].cubin {{.*}}-x-{{.*}}.cubin{{.*}}-y-{{.*}}.cubin{{.*}}-z-{{.*}}.cubin
 
 //
 // Check that the LTO interface works and properly preserves symbols used in a
diff --git a/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp b/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
index b9767a7..bc191af 100644
--- a/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
+++ b/clang/tools/clang-nvlink-wrapper/ClangNVLinkWrapper.cpp
@@ -250,6 +250,7 @@ struct Symbol {
   };
 
   Symbol() : File(), Flags(None), UsedInRegularObj(false) {}
+  Symbol(Symbol::Flags Flags) : File(), Flags(Flags), UsedInRegularObj(true) {}
 
   Symbol(MemoryBufferRef File, const irsymtab::Reader::SymbolRef Sym)
       : File(File), Flags(0), UsedInRegularObj(false) {
@@ -535,6 +536,8 @@ Expected<SmallVector<StringRef>> getInput(const ArgList &Args) {
 
   bool Extracted = true;
   StringMap<Symbol> SymTab;
+  for (auto &Sym : Args.getAllArgValues(OPT_u))
+    SymTab[Sym] = Symbol(Symbol::Undefined);
   SmallVector<std::unique_ptr<MemoryBuffer>> LinkerInput;
   while (Extracted) {
     Extracted = false;
diff --git a/clang/tools/clang-nvlink-wrapper/NVLinkOpts.td b/clang/tools/clang-nvlink-wrapper/NVLinkOpts.td
index a80c593..6de1a25 100644
--- a/clang/tools/clang-nvlink-wrapper/NVLinkOpts.td
+++ b/clang/tools/clang-nvlink-wrapper/NVLinkOpts.td
@@ -43,11 +43,11 @@ def plugin : JoinedOrSeparate<["--", "-"], "plugin">,
   Flags<[HelpHidden, WrapperOnlyOption]>;
 
 def arch : Separate<["--", "-"], "arch">,
-  HelpText<"Specify the 'sm_' name of the target architecture.">;
+  HelpText<"Specify the 'sm_' name of the target architecture">;
 def : Joined<["--", "-"], "plugin-opt=mcpu=">,
   Flags<[HelpHidden, WrapperOnlyOption]>, Alias<arch>;
 
-def g : Flag<["-"], "g">, HelpText<"Specify that this was a debug compile.">;
+def g : Flag<["-"], "g">, HelpText<"Specify that this was a debug compile">;
 def debug : Flag<["--"], "debug">, Alias<g>;
 
 def lto_emit_llvm : Flag<["--"], "lto-emit-llvm">, Flags<[WrapperOnlyOption]>,
@@ -55,6 +55,9 @@ def lto_emit_llvm : Flag<["--"], "lto-emit-llvm">, Flags<[WrapperOnlyOption]>,
 def lto_emit_asm : Flag<["--"], "lto-emit-asm">, Flags<[WrapperOnlyOption]>,
   HelpText<"Emit assembly code">;
 
+def u : JoinedOrSeparate<["-"], "u">, HelpText<"Force undefined symbol during linking">;
+def undefined : JoinedOrSeparate<["--"], "undefined">, Alias<u>;
+
 def O : Joined<["--", "-"], "plugin-opt=O">,
   Flags<[WrapperOnlyOption]>, MetaVarName<"<O0, O1, O2, or O3>">,
   HelpText<"Optimization level for LTO">;