[LinkerWrapper] Only import static libraries with needed symbols

[jhuber6]

Summary

Currently, we pull in every single static archive member as long as we
have an offloading architecture that requires it. This goes against the
standard sematnics of static libraries that only pull in symbols that
define currently undefined symbols. In order to support this we roll
some custom symbol resolution logic to check if a static library is
needed. Because of offloading semantics, this requires an extra check
for externally visibile symbols. E.g. if a static member defines a
kernel we should import it.

The main benefit to this is that we can now link against the
libomptarget.devicertl.a library unconditionally. This removes the
requirement for users to specify LTO on the link command. This will also
allow us to stop using the amdgcn bitcode versions of the libraries.

clang foo.c -fopenmp --offload-arch=gfx1030 -foffload-lto -c
clang foo.o -fopenmp --offload-arch=gfx1030 -foffload-lto

[tra]

@MaskRay - we seem to be reinventing the linker here and could use your expertise.

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

1217–1218

How do we handle conflicting symbols defined more than once?

[jhuber6]

In D142484#4077811, @tra wrote:

@MaskRay - we seem to be reinventing the linker here and could use your expertise.

Yeah, this reinvents a lot of logic. But I don't think there's an easy way to get around this without duplicating some ugly logic. I would absolutely love to shove the vast majority of this LTO / static library handling directly into the linker. E.g. we package the extracted libraries into a new one and let lld handle it. The problem is pretty firmly in Nvidia's hands because they don't have a real linker. They use nvlink which does some magic to .cubin files and doesn't understand our LTO or static library linking. So in order to make this common I needed to reproduce a lot of that logic here.

Also, the one subtle difference in this and the regular rules for static libraries is that we still want to extract symbols that would normally be considered "ExportDynamic". This is just because most of the time, these symbols are going to be read from the user's offloading app.

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

1217–1218

I just assume that will be caught by the actual linker, it's not relevant for knowing which symbols to pull out.

[tra]

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

1217–1218

That would be the case when we were loading all libraries at once.
However, now that we load them as needed one by one, we may end up loading only the first file which provides such symbol and would potentially ignore others, if no other symbol requires loading them. Given that we're checking the symbols anyways, it would be useful to diagnose it, IMO.

[jhuber6]

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

1217–1218

This is the expected behavior of static libraries right? lld has some option to print out why a certain library was extracted, we could do something similar?

[tra]

We could also use more test cases. E.g. weak symbols (should not cause object extraction)

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

1197

This could use some comments.
@MaskRay's blog post https://maskray.me/blog/2021-06-20-symbol-processing#archive-processing would be helpful to explain what's going on here.

I'd also refactor the condition a bit for readability:

bool ResolvesStrongReference = (OldSym & Sym_Undefined) &&
                               !(OldSym & Sym_Weak) &&
                               !(*FlagsOrErr & SymbolRef::SF_Undefined);
bool NewGlobalSymbol = ((NewSymbol || (OldSym & Sym_Undefined)) &&
                        !(*FlagsOrErr & SymbolRef::SF_Undefined) &&
                        !(*FlagsOrErr & SymbolRef::SF_Hidden));
Resolved |= ResolvesStrongReference || NewGlobalSymbol;

1217–1218

Yes, it is expected. It's also a common source of surprises. I'm not sure whether lld diagnoses it these days. We should follow whatever it does.

[jhuber6]

In D142484#4077979, @tra wrote:

We could also use more test cases. E.g. weak symbols (should not cause object extraction)

Yeah, I'll try to add a reasonable test here. It's a little difficult due to the indirect nature of this however.

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

1217–1218

I just tested it with the same library defining a global x for main.c.

> clang main.c libx.a libx.a libx.a -Wl,--why-extract=-
reference	extracted	symbol
/tmp/main-8a5566.o	libx.a(x.o)	x

I could maybe add a similar option here to at least let us know if it extracts. Maybe for some tests.

[jhuber6]

Adding test and making the logic a bit more readable.

[tra]

LGTM. Please wait a bit before landing it, in case @MaskRay has something to say.

clang/test/Driver/linker-wrapper-libs.c

39–40

I'm puzzled about how exactly we're checking the that we extract a static library defining an undefined symbol part here.

Presumably the trailing .o would be the object extracted from the library. I'm not sure where the first .o/.s argument come from? Is that the 'main' GPU object/executable embedded in the host object?

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

1213

Nit: perhaps we should rename Resolved -> ShouldExtractObject, it would be a better match. After all, defining/including a new symbol is not quite the same as resolving a reference.

[jhuber6]

In D142484#4078377, @tra wrote:

LGTM. Please wait a bit before landing it, in case @MaskRay has something to say.

I'm somewhat hoping to get this in before the fork that happens in a few hours.

clang/test/Driver/linker-wrapper-libs.c

39–40

Yes, if we extract it we'll have two object files. The .s comes from the fact that I deliberately made the main binary an ELF while the library is a .bc. This is why we get a .s for Nvidia, it's PTX.

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

1213

Good point, I'll change it.

[tra]

I'm somewhat hoping to get this in before the fork that happens in a few hours.

OK. We can fix it later if it turns out that we've missed something.

[Jake-Egan]

Hi, this new test fails on AIX. Could you take a look?
https://lab.llvm.org/buildbot/#/builders/214/builds/5477/steps/6/logs/FAIL__Clang__linker-wrapper-libs_c

[jhuber6]

In D142484#4079869, @Jake-Egan wrote:

Hi, this new test fails on AIX. Could you take a look?
https://lab.llvm.org/buildbot/#/builders/214/builds/5477/steps/6/logs/FAIL__Clang__linker-wrapper-libs_c

I'm actually not sure why the test it failing this way. The failure suggests that it's not extracting the members, but I have absolutely no clue what could be different on AIX versus all the other systems the tests pass on. Nothing in this logic should be target dependent, unless there's something weird with how we handle llvm-ar or similar.

[jhuber6]

In D142484#4079869, @Jake-Egan wrote:

Hi, this new test fails on AIX. Could you take a look?
https://lab.llvm.org/buildbot/#/builders/214/builds/5477/steps/6/logs/FAIL__Clang__linker-wrapper-libs_c

I might just put that this test requires Linux. Because I don't think I can debug it without access to the system. We don't really support offloading on anything but Linux anyway.

[Jake-Egan]

In D142484#4079909, @jhuber6 wrote:

In D142484#4079869, @Jake-Egan wrote:

Hi, this new test fails on AIX. Could you take a look?
https://lab.llvm.org/buildbot/#/builders/214/builds/5477/steps/6/logs/FAIL__Clang__linker-wrapper-libs_c

I might just put that this test requires Linux. Because I don't think I can debug it without access to the system. We don't really support offloading on anything but Linux anyway.

I think that makes sense. We'd need time to investigate it on AIX anyways.

[DiggerLin]

In D142484#4080202, @Jake-Egan wrote:

In D142484#4079909, @jhuber6 wrote:

In D142484#4079869, @Jake-Egan wrote:

Hi, this new test fails on AIX. Could you take a look?
https://lab.llvm.org/buildbot/#/builders/214/builds/5477/steps/6/logs/FAIL__Clang__linker-wrapper-libs_c

I might just put that this test requires Linux. Because I don't think I can debug it without access to the system. We don't really support offloading on anything but Linux anyway.

I think that makes sense. We'd need time to investigate it on AIX anyways.

I put a patch to fix the problem. https://reviews.llvm.org/D145600

clang/test/Driver/linker-wrapper-libs.c

+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.elf.o
+
+#if defined(RESOLVES)
+int __attribute__((visibility("hidden"))) sym;
+#elif defined(GLOBAL)
+int __attribute__((visibility("protected"))) global;
+#elif defined(WEAK)
+int __attribute__((visibility("hidden"))) weak;
+#elif defined(HIDDEN)
+int __attribute__((visibility("hidden"))) hidden;
+#else
+extern int sym;
+
+extern int __attribute__((weak)) weak;
+
+int foo() { return sym; }
+int bar() { return weak; }
+#endif
+
+//
+// Check that we extract a static library defining an undefined symbol.
+//
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -DRESOLVES -o %t.nvptx.resolves.bc
+// RUN: %clang -cc1 %s -triple amdgcn-amd-amdhsa -emit-llvm-bc -DRESOLVES -o %t.amdgpu.resolves.bc
+// RUN: clang-offload-packager -o %t-lib.out \
+// RUN:   --image=file=%t.nvptx.resolves.bc,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.amdgpu.resolves.bc,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t-lib.out
+// RUN: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu --dry-run \
+// RUN:   --linker-path=/usr/bin/ld -- %t.o %t.a -o a.out 2>&1 \
+// RUN: | FileCheck %s --check-prefix=LIBRARY-RESOLVES
+
+// LIBRARY-RESOLVES: clang{{.*}} -o {{.*}}.img --target=amdgcn-amd-amdhsa -mcpu=gfx1030 {{.*}}.o {{.*}}.o
+// LIBRARY-RESOLVES: clang{{.*}} -o {{.*}}.img --target=nvptx64-nvidia-cuda -march=sm_70 {{.*}}.s {{.*}}.o
+
+//
+// Check that we extract a static library that defines a global visibile to the
+// host.
+//
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -DGLOBAL -o %t.nvptx.global.bc
+// RUN: %clang -cc1 %s -triple amdgcn-amd-amdhsa -emit-llvm-bc -DGLOBAL -o %t.amdgpu.global.bc
+// RUN: clang-offload-packager -o %t-lib.out \
+// RUN:   --image=file=%t.nvptx.global.bc,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.amdgpu.global.bc,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t-lib.out
+// RUN: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu --dry-run \
+// RUN:   --linker-path=/usr/bin/ld -- %t.o %t.a -o a.out 2>&1 \
+// RUN: | FileCheck %s --check-prefix=LIBRARY-GLOBAL
+
+// LIBRARY-GLOBAL: clang{{.*}} -o {{.*}}.img --target=amdgcn-amd-amdhsa -mcpu=gfx1030 {{.*}}.o {{.*}}.o
+// LIBRARY-GLOBAL: clang{{.*}} -o {{.*}}.img --target=nvptx64-nvidia-cuda -march=sm_70 {{.*}}.s {{.*}}.o
+
+//
+// Check that we do not extract an external weak symbol.
+//
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -DWEAK -o %t.nvptx.weak.bc
+// RUN: %clang -cc1 %s -triple amdgcn-amd-amdhsa -emit-llvm-bc -DWEAK -o %t.amdgpu.weak.bc
+// RUN: clang-offload-packager -o %t-lib.out \
+// RUN:   --image=file=%t.nvptx.weak.bc,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.amdgpu.weak.bc,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t-lib.out
+// RUN: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu --dry-run \
+// RUN:   --linker-path=/usr/bin/ld -- %t.o %t.a -o a.out 2>&1 \
+// RUN: | FileCheck %s --check-prefix=LIBRARY-WEAK
+
+// LIBRARY-WEAK: clang{{.*}} -o {{.*}}.img --target=amdgcn-amd-amdhsa -mcpu=gfx1030
+// LIBRARY-WEAK-NOT: {{.*}}.o {{.*}}.o
+// LIBRARY-WEAK: clang{{.*}} -o {{.*}}.img --target=nvptx64-nvidia-cuda -march=sm_70
+
+//
+// Check that we do not extract an unneeded hidden symbol.
+//
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -DHIDDEN -o %t.nvptx.hidden.bc
+// RUN: %clang -cc1 %s -triple amdgcn-amd-amdhsa -emit-llvm-bc -DHIDDEN -o %t.amdgpu.hidden.bc
+// RUN: clang-offload-packager -o %t-lib.out \
+// RUN:   --image=file=%t.nvptx.hidden.bc,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.amdgpu.hidden.bc,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t-lib.out
+// RUN: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu --dry-run \
+// RUN:   --linker-path=/usr/bin/ld -- %t.o %t.a -o a.out 2>&1 \
+// RUN: | FileCheck %s --check-prefix=LIBRARY-HIDDEN
+
+// LIBRARY-HIDDEN: clang{{.*}} -o {{.*}}.img --target=amdgcn-amd-amdhsa -mcpu=gfx1030
+// LIBRARY-HIDDEN-NOT: {{.*}}.o {{.*}}.o
+// LIBRARY-HIDDEN: clang{{.*}} -o {{.*}}.img --target=nvptx64-nvidia-cuda -march=sm_70
+
+//
+// Check that we do not extract a static library that defines a global visibile
+// to the host that is already defined.
+//
+// RUN: %clang -cc1 %s -triple nvptx64-nvidia-cuda -emit-llvm-bc -DGLOBAL -o %t.nvptx.global.bc
+// RUN: %clang -cc1 %s -triple amdgcn-amd-amdhsa -emit-llvm-bc -DGLOBAL -o %t.amdgpu.global.bc
+// RUN: clang-offload-packager -o %t-lib.out \
+// RUN:   --image=file=%t.nvptx.global.bc,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.amdgpu.global.bc,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t-lib.out
+// RUN: llvm-ar rcs %t.a %t.o
+// RUN: clang-offload-packager -o %t.out \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \
+// RUN:   --image=file=%t.elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx1030
+// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out
+// RUN: clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu --dry-run \
+// RUN:   --linker-path=/usr/bin/ld -- %t.o %t.a %t.a -o a.out 2>&1 \
+// RUN: | FileCheck %s --check-prefix=LIBRARY-GLOBAL-DEFINED
+
+// LIBRARY-GLOBAL-DEFINED: clang{{.*}} -o {{.*}}.img --target=amdgcn-amd-amdhsa -mcpu=gfx1030 {{.*}}.o {{.*}}.o
+// LIBRARY-GLOBAL-DEFINED-NOT: {{.*}}.o
+// LIBRARY-GLOBAL-DEFINED: clang{{.*}} -o {{.*}}.img --target=nvptx64-nvidia-cuda -march=sm_70 {{.*}}.s {{.*}}.o

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

 /// `-lfoo` or `-l:libfoo.a`.
 std::optional<std::string> searchLibrary(StringRef Input, StringRef Root,
                                          ArrayRef<StringRef> SearchPaths) {
   if (Input.startswith(":"))
     return findFromSearchPaths(Input.drop_front(), Root, SearchPaths);
   return searchLibraryBaseName(Input, Root, SearchPaths);
 }
 
-/// Search the input files and libraries for embedded device offloading code and
-/// add it to the list of files to be linked. Files coming from static libraries
-/// are only added to the input if they are used by an existing input file.
+/// Common redeclaration of needed symbol flags.
+enum Symbol : uint32_t {
+  Sym_None = 0,
+  Sym_Undefined = 1U << 1,
+  Sym_Weak = 1U << 2,
+};
+
+/// Scan the symbols from a BitcodeFile \p Buffer and record if we need to
+/// extract any symbols from it.
+Expected<bool> getSymbolsFromBitcode(MemoryBufferRef Buffer, StringSaver &Saver,
+                                     DenseMap<StringRef, Symbol> &Syms) {
+  Expected<IRSymtabFile> IRSymtabOrErr = readIRSymtab(Buffer);
+  if (!IRSymtabOrErr)
+    return IRSymtabOrErr.takeError();
+
+  bool ShouldExtract = false;
+  for (unsigned I = 0; I != IRSymtabOrErr->Mods.size(); ++I) {
+    for (const auto &Sym : IRSymtabOrErr->TheReader.module_symbols(I)) {
+      if (Sym.isFormatSpecific() || !Sym.isGlobal())
+        continue;
+
+      bool NewSymbol = Syms.count(Sym.getName()) == 0;
+      auto &OldSym = Syms[Saver.save(Sym.getName())];
+
+      // We will extract if it defines a currenlty undefined non-weak symbol.
+      bool ResolvesStrongReference =
+          ((OldSym & Sym_Undefined && !(OldSym & Sym_Weak)) &&
+           !Sym.isUndefined());
+      // We will extract if it defines a new global symbol visible to the host.
+      bool NewGlobalSymbol =
+          ((NewSymbol || (OldSym & Sym_Undefined)) && !Sym.isUndefined() &&
+           !Sym.canBeOmittedFromSymbolTable() &&
+           (Sym.getVisibility() != GlobalValue::HiddenVisibility));
+      ShouldExtract |= ResolvesStrongReference | NewGlobalSymbol;
+
+      // Update this symbol in the "table" with the new information.
+      if (OldSym & Sym_Undefined && !Sym.isUndefined())
+        OldSym = static_cast<Symbol>(OldSym & ~Sym_Undefined);
+      if (Sym.isUndefined() && NewSymbol)
+        OldSym = static_cast<Symbol>(OldSym | Sym_Undefined);
+      if (Sym.isWeak())
+        OldSym = static_cast<Symbol>(OldSym | Sym_Weak);
+    }
+  }
+
+  return ShouldExtract;
+}
+
+/// Scan the symbols from an ObjectFile \p Obj and record if we need to extract
+/// any symbols from it.
+Expected<bool> getSymbolsFromObject(const ObjectFile &Obj, StringSaver &Saver,
+                                    DenseMap<StringRef, Symbol> &Syms) {
+  bool ShouldExtract = false;
+  for (SymbolRef Sym : Obj.symbols()) {
+    auto FlagsOrErr = Sym.getFlags();
+    if (!FlagsOrErr)
+      return FlagsOrErr.takeError();
+
+    if (!(*FlagsOrErr & SymbolRef::SF_Global) ||
+        (*FlagsOrErr & SymbolRef::SF_FormatSpecific))
+      continue;
+
+    auto NameOrErr = Sym.getName();
+    if (!NameOrErr)
+      return NameOrErr.takeError();
+
+    bool NewSymbol = Syms.count(*NameOrErr) == 0;
+    auto &OldSym = Syms[Saver.save(*NameOrErr)];
+
+    // We will extract if it defines a currenlty undefined non-weak symbol.
+    bool ResolvesStrongReference = (OldSym & Sym_Undefined) &&
+                                   !(OldSym & Sym_Weak) &&
+                                   !(*FlagsOrErr & SymbolRef::SF_Undefined);
+
+    // We will extract if it defines a new global symbol visible to the host.
+    bool NewGlobalSymbol = ((NewSymbol || (OldSym & Sym_Undefined)) &&
+                            !(*FlagsOrErr & SymbolRef::SF_Undefined) &&
+                            !(*FlagsOrErr & SymbolRef::SF_Hidden));
+    ShouldExtract |= ResolvesStrongReference | NewGlobalSymbol;
+
+    // Update this symbol in the "table" with the new information.
+    if (OldSym & Sym_Undefined && !(*FlagsOrErr & SymbolRef::SF_Undefined))
+      OldSym = static_cast<Symbol>(OldSym & ~Sym_Undefined);
+    if (*FlagsOrErr & SymbolRef::SF_Undefined && NewSymbol)
+      OldSym = static_cast<Symbol>(OldSym | Sym_Undefined);
+    if (*FlagsOrErr & SymbolRef::SF_Weak)
+      OldSym = static_cast<Symbol>(OldSym | Sym_Weak);
+  }
+  return ShouldExtract;
+}
+
+/// Attempt to 'resolve' symbols found in input files. We use this to
+/// determine if an archive member needs to be extracted. An archive member
+/// will be extracted if any of the following is true.
+///   1) It defines an undefined symbol in a regular object filie.
+///   2) It defines a global symbol without hidden visibility that has not
+///      yet been defined.
+Expected<bool> getSymbols(StringRef Image, StringSaver &Saver,
+                          DenseMap<StringRef, Symbol> &Syms) {
+  MemoryBufferRef Buffer = MemoryBufferRef(Image, "");
+  switch (identify_magic(Image)) {
+  case file_magic::bitcode:
+    return getSymbolsFromBitcode(Buffer, Saver, Syms);
+  case file_magic::elf_relocatable: {
+    Expected<std::unique_ptr<ObjectFile>> ObjFile =
+        ObjectFile::createObjectFile(Buffer);
+    if (!ObjFile)
+      return ObjFile.takeError();
+    return getSymbolsFromObject(**ObjFile, Saver, Syms);
+  }
+  default:
+    return false;
+  }
+}
+
+/// Search the input files and libraries for embedded device offloading code
+/// and add it to the list of files to be linked. Files coming from static
+/// libraries are only added to the input if they are used by an existing
+/// input file.
 Expected<SmallVector<OffloadFile>> getDeviceInput(const ArgList &Args) {
   llvm::TimeTraceScope TimeScope("ExtractDeviceCode");
 
   StringRef Root = Args.getLastArgValue(OPT_sysroot_EQ);
   SmallVector<StringRef> LibraryPaths;
   for (const opt::Arg *Arg : Args.filtered(OPT_library_path))
     LibraryPaths.push_back(Arg->getValue());
 
+  BumpPtrAllocator Alloc;
+  StringSaver Saver(Alloc);
+
   // Try to extract device code from the linker input files.
   SmallVector<OffloadFile> InputFiles;
-  SmallVector<OffloadFile> LazyInputFiles;
-  for (const opt::Arg *Arg : Args.filtered(OPT_INPUT)) {
-    StringRef Filename = Arg->getValue();
-    if (!sys::fs::exists(Filename) || sys::fs::is_directory(Filename))
+  DenseMap<OffloadFile::TargetID, DenseMap<StringRef, Symbol>> Syms;
+  for (const opt::Arg *Arg : Args.filtered(OPT_INPUT, OPT_library)) {
+    std::optional<std::string> Filename =
+        Arg->getOption().matches(OPT_library)
+            ? searchLibrary(Arg->getValue(), Root, LibraryPaths)
+            : std::string(Arg->getValue());
+
+    if (!Filename && Arg->getOption().matches(OPT_library))
+      reportError(createStringError(inconvertibleErrorCode(),
+                                    "unable to find library -l%s",
+                                    Arg->getValue()));
+
+    if (!Filename || !sys::fs::exists(*Filename) ||
+        sys::fs::is_directory(*Filename))
       continue;
 
     ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
-        MemoryBuffer::getFileOrSTDIN(Filename);
+        MemoryBuffer::getFileOrSTDIN(*Filename);
     if (std::error_code EC = BufferOrErr.getError())
-      return createFileError(Filename, EC);
+      return createFileError(*Filename, EC);
 
-    if (identify_magic((*BufferOrErr)->getBuffer()) ==
-        file_magic::elf_shared_object)
+    MemoryBufferRef Buffer = **BufferOrErr;
+    if (identify_magic(Buffer.getBuffer()) == file_magic::elf_shared_object)
       continue;
 
-    bool IsLazy =
-        identify_magic((*BufferOrErr)->getBuffer()) == file_magic::archive;
-    if (Error Err = extractOffloadBinaries(
-            **BufferOrErr, IsLazy ? LazyInputFiles : InputFiles))
+    SmallVector<OffloadFile> Binaries;
+    if (Error Err = extractOffloadBinaries(Buffer, Binaries))
       return std::move(Err);
-  }
 
-  // Try to extract input from input archive libraries.
-  for (const opt::Arg *Arg : Args.filtered(OPT_library)) {
-    if (auto Library = searchLibrary(Arg->getValue(), Root, LibraryPaths)) {
-      ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
-          MemoryBuffer::getFileOrSTDIN(*Library);
-      if (std::error_code EC = BufferOrErr.getError())
-        reportError(createFileError(*Library, EC));
+    // We only extract archive members that are needed.
+    bool IsArchive = identify_magic(Buffer.getBuffer()) == file_magic::archive;
+    bool Extracted = true;
+    while (Extracted) {
+      Extracted = false;
+      for (OffloadFile &Binary : Binaries) {
+        if (!Binary.getBinary())
+          continue;
 
-      if (identify_magic((*BufferOrErr)->getBuffer()) != file_magic::archive)
+        // If we don't have an object file for this architecture do not
+        // extract.
+        if (IsArchive && !Syms.count(Binary))
           continue;
 
-      if (Error Err = extractOffloadBinaries(**BufferOrErr, LazyInputFiles))
-        return std::move(Err);
-    } else {
-      reportError(createStringError(inconvertibleErrorCode(),
-                                    "unable to find library -l%s",
-                                    Arg->getValue()));
+        Expected<bool> ExtractOrErr =
+            getSymbols(Binary.getBinary()->getImage(), Saver, Syms[Binary]);
+        if (!ExtractOrErr)
+          return ExtractOrErr.takeError();
+
+        Extracted = IsArchive && *ExtractOrErr;
+
+        if (!IsArchive || Extracted)
+          InputFiles.emplace_back(std::move(Binary));
+
+        // If we extracted any files we need to check all the symbols again.
+        if (Extracted)
+          break;
+      }
     }
   }
 
   for (StringRef Library : Args.getAllArgValues(OPT_bitcode_library_EQ)) {
     auto FileOrErr = getInputBitcodeLibrary(Library);
     if (!FileOrErr)
       return FileOrErr.takeError();
     InputFiles.push_back(std::move(*FileOrErr));
   }
 
-  DenseSet<OffloadFile::TargetID> IsTargetUsed;
-  for (const auto &File : InputFiles)
-    IsTargetUsed.insert(File);
-
-  // We should only include input files that are used.
-  // TODO: Only load a library if it defined undefined symbols in the input.
-  for (auto &LazyFile : LazyInputFiles)
-    if (IsTargetUsed.contains(LazyFile))
-      InputFiles.emplace_back(std::move(LazyFile));
-
   return std::move(InputFiles);
 }
 
 } // namespace
 
 int main(int Argc, char **Argv) {
   InitLLVM X(Argc, Argv);
   InitializeAllTargetInfos();