[LTO API] add lto_codegen_set_module to set the destination module.

When debugging LTO issues with ld64, we use -save-temps to save the merged
optimized bitcode file, then invoke ld64 again on the single bitcode file to
speed up debugging code generation passes and ld64 stuff after code generation.

llvm linking a single bitcode file via lto_codegen_add_module will generate a
different bitcode file from the single input. With the newly-added
lto_codegen_set_module, we can make sure the destination module is the same as
the input.

lto_codegen_set_module will transfer the ownship of the module to code
generator.

rdar://19024554


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@230290 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm-c/lto.h

@@ -40,7 +40,7 @@ typedef bool lto_bool_t;
  * @{
  */
 
-#define LTO_API_VERSION 12
+#define LTO_API_VERSION 13
 
 /**
  * \since prior to LTO_API_VERSION=3
@@ -395,6 +395,17 @@ lto_codegen_dispose(lto_code_gen_t);
 extern lto_bool_t
 lto_codegen_add_module(lto_code_gen_t cg, lto_module_t mod);
 
+/**
+ * Sets the object module for code generation. This will transfer the ownship of
+ * the module to code generator.
+ *
+ * \c cg and \c mod must both be in the same context.
+ *
+ * \since prior to LTO_API_VERSION=13
+ */
+extern void
+lto_codegen_set_module(lto_code_gen_t cg, lto_module_t mod);
+
 /**
  * Sets if debug info should be generated.
  * Returns true on error (check lto_get_error_message() for details).

include/llvm/LTO/LTOCodeGenerator.h

@@ -67,6 +67,9 @@ struct LTOCodeGenerator {
   // Merge given module, return true on success.
   bool addModule(struct LTOModule *);
 
+  // Set the destination module.
+  void setModule(struct LTOModule *);
+
   void setTargetOptions(TargetOptions options);
   void setDebugInfo(lto_debug_model);
   void setCodePICModel(lto_codegen_model);

include/llvm/Linker/Linker.h

@@ -70,6 +70,9 @@ public:
   /// Returns true on error.
   bool linkInModule(Module *Src);
 
+  /// \brief Set the composite to the passed-in module.
+  void setModule(Module *Dst);
+
   static bool LinkModules(Module *Dest, Module *Src,
                           DiagnosticHandlerFunction DiagnosticHandler);
 

lib/LTO/LTOCodeGenerator.cpp

@@ -141,6 +141,22 @@ bool LTOCodeGenerator::addModule(LTOModule *mod) {
   return !ret;
 }
 
+void LTOCodeGenerator::setModule(LTOModule *Mod) {
+  assert(&Mod->getModule().getContext() == &Context &&
+         "Expected module in same context");
+
+  // Delete the old merged module.
+  if (IRLinker.getModule())
+    IRLinker.deleteModule();
+  AsmUndefinedRefs.clear();
+
+  IRLinker.setModule(&Mod->getModule());
+
+  const std::vector<const char*> &Undefs = Mod->getAsmUndefinedRefs();
+  for (int I = 0, E = Undefs.size(); I != E; ++I)
+    AsmUndefinedRefs[Undefs[I]] = 1;
+}
+
 void LTOCodeGenerator::setTargetOptions(TargetOptions options) {
   Options = options;
 }

lib/Linker/LinkModules.cpp

@@ -1748,6 +1748,10 @@ bool Linker::linkInModule(Module *Src) {
   return RetCode;
 }
 
+void Linker::setModule(Module *Dst) {
+  init(Dst, DiagnosticHandler);
+}
+
 //===----------------------------------------------------------------------===//
 // LinkModules entrypoint.
 //===----------------------------------------------------------------------===//

test/LTO/X86/set-merged.ll

@@ -0,0 +1,36 @@
+; RUN: llvm-as < %s >%t1
+; RUN: llvm-lto -exported-symbol=_main -set-merged-module -o %t2 %t1
+; RUN: llvm-objdump -d %t2 | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.10.0"
+
+; CHECK: _main
+; CHECK: movl $132
+define i32 @_Z3fooi(i32 %a) {
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %0 = load i32* %a.addr, align 4
+  %1 = load i32* %a.addr, align 4
+  %call = call i32 @_Z4bar2i(i32 %1)
+  %add = add nsw i32 %0, %call
+  ret i32 %add
+}
+
+define i32 @_Z4bar2i(i32 %a) {
+entry:
+  %a.addr = alloca i32, align 4
+  store i32 %a, i32* %a.addr, align 4
+  %0 = load i32* %a.addr, align 4
+  %mul = mul nsw i32 2, %0
+  ret i32 %mul
+}
+
+define i32 @main() {
+entry:
+  %retval = alloca i32, align 4
+  store i32 0, i32* %retval
+  %call = call i32 @_Z3fooi(i32 44)
+  ret i32 %call
+}

tools/llvm-lto/llvm-lto.cpp

@@ -69,6 +69,10 @@ static cl::opt<bool> ListSymbolsOnly(
     "list-symbols-only", cl::init(false),
     cl::desc("Instead of running LTO, list the symbols in each IR file"));
 
+static cl::opt<bool> SetMergedModule(
+    "set-merged-module", cl::init(false),
+    cl::desc("Use the first input module as the merged module"));
+
 namespace {
 struct ModuleInfo {
   std::vector<bool> CanBeHidden;
@@ -194,15 +198,22 @@ int main(int argc, char **argv) {
       return 1;
     }
 
-    if (!CodeGen.addModule(Module.get()))
+    LTOModule *LTOMod = Module.get();
+
+    // We use the first input module as the destination module when
+    // SetMergedModule is true.
+    if (SetMergedModule && i == BaseArg) {
+      // Transfer ownership to the code generator.
+      CodeGen.setModule(Module.release());
+    } else if (!CodeGen.addModule(Module.get()))
       return 1;
 
-    unsigned NumSyms = Module->getSymbolCount();
+    unsigned NumSyms = LTOMod->getSymbolCount();
     for (unsigned I = 0; I < NumSyms; ++I) {
-      StringRef Name = Module->getSymbolName(I);
+      StringRef Name = LTOMod->getSymbolName(I);
       if (!DSOSymbolsSet.count(Name))
         continue;
-      lto_symbol_attributes Attrs = Module->getSymbolAttributes(I);
+      lto_symbol_attributes Attrs = LTOMod->getSymbolAttributes(I);
       unsigned Scope = Attrs & LTO_SYMBOL_SCOPE_MASK;
       if (Scope != LTO_SYMBOL_SCOPE_DEFAULT_CAN_BE_HIDDEN)
         KeptDSOSyms.push_back(Name);

tools/lto/lto.cpp

@@ -248,6 +248,10 @@ bool lto_codegen_add_module(lto_code_gen_t cg, lto_module_t mod) {
   return !unwrap(cg)->addModule(unwrap(mod));
 }
 
+void lto_codegen_set_module(lto_code_gen_t cg, lto_module_t mod) {
+  unwrap(cg)->setModule(unwrap(mod));
+}
+
 bool lto_codegen_set_debug_model(lto_code_gen_t cg, lto_debug_model debug) {
   unwrap(cg)->setDebugInfo(debug);
   return false;

tools/lto/lto.exports

@@ -25,6 +25,7 @@ lto_module_dispose
 lto_api_version
 lto_codegen_set_diagnostic_handler
 lto_codegen_add_module
+lto_codegen_set_module
 lto_codegen_add_must_preserve_symbol
 lto_codegen_compile
 lto_codegen_create