Implement selective bitcode linking and internalization of linked symbols

[tra]

Summary

In order to support CUDA we need a way to link in subset of functions from bitcode library.

Current way to do that suggested in NVPTX guide (http://llvm.org/docs/NVPTXUsage.html) requires linking in complete library, internalizing all symbols except those that were originally present in TU before linking and running GDCE pass to eliminate bitcode we don't need.

Considering that we only need fairly small subset of functions from the library, better way to do that would be to directly link in only the symbols needed by the destination module and internalize them in process, if required.

This patch adds two new linker flags to do exactly that.

• -only-needed -- links in only symbols needed by destination module
• -internalize -- internalize linked symbols.

[tra]

lib/Target/NVPTX/NVPTXTargetMachine.cpp

184–186

This change was included by mistake. I'll update the patch shortly to remove it.

[tra]

Removed changes included by mistake.

[tra]

Fixed a nit spotted by dexonsmith@.

[tra]

LGTM by dexonsmith@ on llvm-commits list.

include/llvm/Linker/Linker.h

struct IdentifiedStructTypeSet {
 
     void addNonOpaque(StructType *Ty);
     void switchToNonOpaque(StructType *Ty);
     void addOpaque(StructType *Ty);
     StructType *findNonOpaque(ArrayRef<Type *> ETypes, bool IsPacked);
     bool hasType(StructType *Ty);
   };
 
+  enum Flags {
+    None = 0,
+    OverrideFromSrc = (1 << 0),
+    LinkOnlyNeeded = (1 << 1),
+    InternalizeLinkedSymbols = (1 << 2)
+  };
+
   Linker(Module *M, DiagnosticHandlerFunction DiagnosticHandler);
   Linker(Module *M);
 
   Module *getModule() const { return Composite; }
   void deleteModule();
 
   /// \brief Link \p Src into the composite. The source is destroyed.
   /// Passing OverrideSymbols as true will have symbols from Src
   /// shadow those in the Dest.
   /// Returns true on error.
-  bool linkInModule(Module *Src, bool OverrideSymbols = false);
+  bool linkInModule(Module *Src, unsigned Flags = Flags::None);
 
   /// \brief Set the composite to the passed-in module.
   void setModule(Module *Dst);
 
   static bool LinkModules(Module *Dest, Module *Src,
-                          DiagnosticHandlerFunction DiagnosticHandler);
+                          DiagnosticHandlerFunction DiagnosticHandler,
+                          unsigned Flags = Flags::None);
 
-  static bool LinkModules(Module *Dest, Module *Src);
+  static bool LinkModules(Module *Dest, Module *Src,
+                          unsigned Flags = Flags::None);
 
 private:
   void init(Module *M, DiagnosticHandlerFunction DiagnosticHandler);
   Module *Composite;
 
   IdentifiedStructTypeSet IdentifiedStructTypes;
 
   DiagnosticHandlerFunction DiagnosticHandler;
 };
 
 } // End llvm namespace
 
 #endif

lib/Linker/LinkModules.cpp

class ModuleLinker {
   std::vector<GlobalValue *> LazilyLinkGlobalValues;
 
   /// Functions that have replaced other functions.
   SmallPtrSet<const Function *, 16> OverridingFunctions;
 
   DiagnosticHandlerFunction DiagnosticHandler;
 
   /// For symbol clashes, prefer those from Src.
-  bool OverrideFromSrc;
+  unsigned Flags;
 
 public:
   ModuleLinker(Module *dstM, Linker::IdentifiedStructTypeSet &Set, Module *srcM,
-               DiagnosticHandlerFunction DiagnosticHandler,
-               bool OverrideFromSrc)
+               DiagnosticHandlerFunction DiagnosticHandler, unsigned Flags)
       : DstM(dstM), SrcM(srcM), TypeMap(Set),
         ValMaterializer(TypeMap, DstM, LazilyLinkGlobalValues),
-        DiagnosticHandler(DiagnosticHandler), OverrideFromSrc(OverrideFromSrc) {
-  }
+        DiagnosticHandler(DiagnosticHandler), Flags(Flags) {}
 
   bool run();
 
+  bool shouldOverrideFromSrc() { return Flags & Linker::OverrideFromSrc; }
+  bool shouldLinkOnlyNeeded() { return Flags & Linker::LinkOnlyNeeded; }
+  bool shouldInternalizeLinkedSymbols() {
+    return Flags & Linker::InternalizeLinkedSymbols;
+  }
+
 private:
   bool shouldLinkFromSource(bool &LinkFromSrc, const GlobalValue &Dest,
                             const GlobalValue &Src);
 
   /// Helper method for setting a message and returning an error code.
   bool emitError(const Twine &Message) {
     DiagnosticHandler(LinkDiagnosticInfo(DS_Error, Message));
     return true;
bool ModuleLinker::getComdatResult(const Comdat *SrcC,
   return computeResultingSelectionKind(ComdatName, SSK, DSK, Result,
                                        LinkFromSrc);
 }
 
 bool ModuleLinker::shouldLinkFromSource(bool &LinkFromSrc,
                                         const GlobalValue &Dest,
                                         const GlobalValue &Src) {
   // Should we unconditionally use the Src?
-  if (OverrideFromSrc) {
+  if (shouldOverrideFromSrc()) {
     LinkFromSrc = true;
     return false;
   }
 
   // We always have to add Src if it has appending linkage.
   if (Src.hasAppendingLinkage()) {
     LinkFromSrc = true;
     return false;
if (!LinkFromSrc && !DGV)
     return false;
 
   GlobalValue *NewGV;
   if (!LinkFromSrc) {
     NewGV = DGV;
   } else {
     // If the GV is to be lazily linked, don't create it just yet.
     // The ValueMaterializerTy will deal with creating it if it's used.
-    if (!DGV && !OverrideFromSrc &&
+    if (!DGV && !shouldOverrideFromSrc() &&
         (SGV->hasLocalLinkage() || SGV->hasLinkOnceLinkage() ||
          SGV->hasAvailableExternallyLinkage())) {
       DoNotLinkFromSource.insert(SGV);
       return false;
     }
 
+    // When we only want to link in unresolved dependencies, blacklist
+    // the symbol unless unless DestM has a matching declaration (DGV).
+    if (shouldLinkOnlyNeeded() && !(DGV && DGV->isDeclaration())) {
+      DoNotLinkFromSource.insert(SGV);
+      return false;
+    }
+
     NewGV = copyGlobalValueProto(TypeMap, *DstM, SGV);
 
     if (DGV && isa<Function>(DGV))
       if (auto *NewF = dyn_cast<Function>(NewGV))
         OverridingFunctions.insert(NewF);
   }
 
   NewGV->setUnnamedAddr(HasUnnamedAddr);
void ModuleLinker::linkAliasBody(GlobalAlias &Dst, GlobalAlias &Src) {
   Constant *Val = MapValue(Aliasee, ValueMap, RF_MoveDistinctMDs, &TypeMap,
                            &ValMaterializer);
   Dst.setAliasee(Val);
 }
 
 bool ModuleLinker::linkGlobalValueBody(GlobalValue &Src) {
   Value *Dst = ValueMap[&Src];
   assert(Dst);
+  if (shouldInternalizeLinkedSymbols())
+    if (auto *DGV = dyn_cast<GlobalValue>(Dst)) {
+      DGV->setVisibility(GlobalValue::DefaultVisibility);
+      DGV->setLinkage(GlobalValue::InternalLinkage);
+    }
   if (auto *F = dyn_cast<Function>(&Src))
     return linkFunctionBody(cast<Function>(*Dst), *F);
   if (auto *GVar = dyn_cast<GlobalVariable>(&Src)) {
     linkGlobalInit(cast<GlobalVariable>(*Dst), *GVar);
     return false;
   }
   linkAliasBody(cast<GlobalAlias>(*Dst), cast<GlobalAlias>(Src));
   return false;
for (GlobalVariable &Src : SrcM->globals()) {
     linkGlobalValueBody(Src);
   }
 
   // Process vector of lazily linked in functions.
   while (!LazilyLinkGlobalValues.empty()) {
     GlobalValue *SGV = LazilyLinkGlobalValues.back();
     LazilyLinkGlobalValues.pop_back();
 
+    // Skip declarations that ValueMaterializer may have created in
+    // case we link in only some of SrcM.
+    if (shouldLinkOnlyNeeded() && SGV->isDeclaration())
+      continue;
+
     assert(!SGV->isDeclaration() && "users should not pass down decls");
     if (linkGlobalValueBody(*SGV))
       return true;
   }
 
   return false;
 }
 
Linker::Linker(Module *M) {
   });
 }
 
 void Linker::deleteModule() {
   delete Composite;
   Composite = nullptr;
 }
 
-bool Linker::linkInModule(Module *Src, bool OverrideSymbols) {
+bool Linker::linkInModule(Module *Src, unsigned Flags) {
   ModuleLinker TheLinker(Composite, IdentifiedStructTypes, Src,
-                         DiagnosticHandler, OverrideSymbols);
+                         DiagnosticHandler, Flags);
   bool RetCode = TheLinker.run();
   Composite->dropTriviallyDeadConstantArrays();
   return RetCode;
 }
 
 void Linker::setModule(Module *Dst) {
   init(Dst, DiagnosticHandler);
 }
 
 //===----------------------------------------------------------------------===//
 // LinkModules entrypoint.
 //===----------------------------------------------------------------------===//
 
 /// This function links two modules together, with the resulting Dest module
 /// modified to be the composite of the two input modules. If an error occurs,
 /// true is returned and ErrorMsg (if not null) is set to indicate the problem.
 /// Upon failure, the Dest module could be in a modified state, and shouldn't be
 /// relied on to be consistent.
 bool Linker::LinkModules(Module *Dest, Module *Src,
-                         DiagnosticHandlerFunction DiagnosticHandler) {
+                         DiagnosticHandlerFunction DiagnosticHandler,
+                         unsigned Flags) {
   Linker L(Dest, DiagnosticHandler);
-  return L.linkInModule(Src);
+  return L.linkInModule(Src, Flags);
 }
 
-bool Linker::LinkModules(Module *Dest, Module *Src) {
+bool Linker::LinkModules(Module *Dest, Module *Src, unsigned Flags) {
   Linker L(Dest);
-  return L.linkInModule(Src);
+  return L.linkInModule(Src, Flags);
 }
 
 //===----------------------------------------------------------------------===//
 // C API.
 //===----------------------------------------------------------------------===//
 
 LLVMBool LLVMLinkModules(LLVMModuleRef Dest, LLVMModuleRef Src,
                          LLVMLinkerMode Unused, char **OutMessages) {

lib/Target/NVPTX/NVPTXTargetMachine.cpp

extern "C" void LLVMInitializeNVPTXTarget() {
   initializeNVPTXAllocaHoistingPass(PR);
   initializeNVPTXAssignValidGlobalNamesPass(PR);
   initializeNVPTXFavorNonGenericAddrSpacesPass(PR);
   initializeNVPTXLowerKernelArgsPass(PR);
   initializeNVPTXLowerAllocaPass(PR);
   initializeNVPTXLowerAggrCopiesPass(PR);
 }
 
+static cl::opt<bool> NVPTXEnableReflect(
+    "nvptx-enable-reflect", cl::init(false), cl::Hidden,
+    cl::desc("Run NVVM reflection pass before other standard passes, "
+             "disabled by default"));
+
 static std::string computeDataLayout(bool is64Bit) {
   std::string Ret = "e";
 
   if (!is64Bit)
     Ret += "-p:32:32";
 
   Ret += "-i64:64-v16:16-v32:32-n16:32:64";
 
void NVPTXPassConfig::addIRPasses() {
   // around after register allocation (which in our case, is *all* registers).
   // We explicitly disable them here.  We do, however, need some functionality
   // of the PrologEpilogCodeInserter pass, so we emulate that behavior in the
   // NVPTXPrologEpilog pass (see NVPTXPrologEpilogPass.cpp).
   disablePass(&PrologEpilogCodeInserterID);
   disablePass(&MachineCopyPropagationID);
   disablePass(&TailDuplicateID);
 
+  if (NVPTXEnableReflect)
+    addPass(createNVVMReflectPass());
   addPass(createNVPTXImageOptimizerPass());
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
 
   // === Propagate special address spaces ===
   addPass(createNVPTXLowerKernelArgsPass(&getNVPTXTargetMachine()));
   // NVPTXLowerKernelArgs emits alloca for byval parameters which can often
   // be eliminated by SROA.
   addPass(createSROAPass());

test/CodeGen/NVPTX/nvptx-enable-reflect.ll

+; RUN: llc < %s -march=nvptx -debug-pass=Structure -o /dev/null 2>&1 \
+; RUN:    | FileCheck %s --check-prefix=ALL --check-prefix=NOREFLECT
+; RUN: llc < %s -march=nvptx -debug-pass=Structure -o /dev/null -nvptx-enable-reflect 2>&1 \
+; RUN:    | FileCheck %s --check-prefix=ALL --check-prefix=REFLECT
+
+; ALL: Pass Arguments
+; ALL:   ModulePass Manager
+; REFLECT-NEXT: Replace occurrences of __nvvm_reflect
+; NOREFLECT-NOT: Replace occurrences of __nvvm_reflect
+; ALL-NEXT: FunctionPass Manager

test/Linker/Inputs/linkage.c.ll

+@X = global i32 5
+@U = global i32 6
+define i32 @foo() { ret i32 7 }
+define i32 @unused() { ret i32 8 }

test/Linker/link-flags.ll

+; RUN: llvm-as %S/Inputs/linkage.b.ll -o %t.b.bc
+; RUN: llvm-as %S/Inputs/linkage.c.ll -o %t.c.bc
+; RUN: llvm-link -S %t.b.bc %t.c.bc | FileCheck %s -check-prefix=B -check-prefix=C -check-prefix=CU
+; RUN: llvm-link -S -only-needed %t.b.bc %t.c.bc | FileCheck %s -check-prefix=B -check-prefix=C -check-prefix=CN
+; RUN: llvm-link -S -internalize %t.b.bc %t.c.bc | FileCheck %s -check-prefix=B -check-prefix=CI
+; RUN: llvm-link -S -internalize -only-needed %t.b.bc %t.c.bc | FileCheck %s -check-prefix=B -check-prefix=CN
+
+C-LABEL: @X = global i32 5
+CI-LABEL: @X = internal global i32 5
+CU-LABEL:@U = global i32 6
+CI-LABEL:@U = internal global i32 6
+CN-LABEL-NOT:@U
+
+B-LABEL: define void @bar() {
+
+C-LABEL: define i32 @foo()
+CI-LABEL: define internal i32 @foo()
+
+CU-LABEL:define i32 @unused() {
+CI-LABEL:define internal i32 @unused() {
+CN-LABEL-NOT:@unused()

tools/llvm-link/llvm-link.cpp

static cl::list<std::string> OverridingInputs(
     cl::desc(
         "input bitcode file which can override previously defined symbol(s)"));
 
 static cl::opt<std::string>
 OutputFilename("o", cl::desc("Override output filename"), cl::init("-"),
                cl::value_desc("filename"));
 
 static cl::opt<bool>
+Internalize("internalize", cl::desc("Internalize linked symbols"));
+
+static cl::opt<bool>
+OnlyNeeded("only-needed", cl::desc("Link only needed symbols"));
+
+static cl::opt<bool>
 Force("f", cl::desc("Enable binary output on terminals"));
 
 static cl::opt<bool>
 OutputAssembly("S",
          cl::desc("Write output as LLVM assembly"), cl::Hidden);
 
 static cl::opt<bool>
 Verbose("v", cl::desc("Print information about actions taken"));
static void diagnosticHandler(const DiagnosticInfo &DI) {
 
   DiagnosticPrinterRawOStream DP(errs());
   DI.print(DP);
   errs() << '\n';
 }
 
 static bool linkFiles(const char *argv0, LLVMContext &Context, Linker &L,
                       const cl::list<std::string> &Files,
-                      bool OverrideDuplicateSymbols) {
+                      unsigned Flags) {
+  // Filter out flags that don't apply to the first file we load.
+  unsigned ApplicableFlags = Flags & Linker::Flags::OverrideFromSrc;
   for (const auto &File : Files) {
     std::unique_ptr<Module> M = loadFile(argv0, File, Context);
     if (!M.get()) {
       errs() << argv0 << ": error loading file '" << File << "'\n";
       return false;
     }
 
     if (verifyModule(*M, &errs())) {
       errs() << argv0 << ": " << File << ": error: input module is broken!\n";
       return false;
     }
 
     if (Verbose)
       errs() << "Linking in '" << File << "'\n";
 
-    if (L.linkInModule(M.get(), OverrideDuplicateSymbols))
+    if (L.linkInModule(M.get(), ApplicableFlags))
       return false;
+    // All linker flags apply to linking of subsequent files.
+    ApplicableFlags = Flags;
   }
 
   return true;
 }
 
 int main(int argc, char **argv) {
   // Print a stack trace if we signal out.
   sys::PrintStackTraceOnErrorSignal();
   PrettyStackTraceProgram X(argc, argv);
 
   LLVMContext &Context = getGlobalContext();
   llvm_shutdown_obj Y;  // Call llvm_shutdown() on exit.
   cl::ParseCommandLineOptions(argc, argv, "llvm linker\n");
 
   auto Composite = make_unique<Module>("llvm-link", Context);
   Linker L(Composite.get(), diagnosticHandler);
 
+  unsigned Flags = Linker::Flags::None;
+  if (Internalize)
+    Flags |= Linker::Flags::InternalizeLinkedSymbols;
+  if (OnlyNeeded)
+    Flags |= Linker::Flags::LinkOnlyNeeded;
+
   // First add all the regular input files
-  if (!linkFiles(argv[0], Context, L, InputFilenames, false))
+  if (!linkFiles(argv[0], Context, L, InputFilenames, Flags))
     return 1;
 
   // Next the -override ones.
-  if (!linkFiles(argv[0], Context, L, OverridingInputs, true))
+  if (!linkFiles(argv[0], Context, L, OverridingInputs,
+                 Flags | Linker::Flags::OverrideFromSrc))
     return 1;
 
   if (DumpAsm) errs() << "Here's the assembly:\n" << *Composite;
 
   std::error_code EC;
   tool_output_file Out(OutputFilename, EC, sys::fs::F_None);
   if (EC) {
     errs() << EC.message() << '\n';