Add a libLTO API to stop/restart ThinLTO between optimizations and CodeGen

This allows the linker to instruct ThinLTO to perform only the
optimization part or only the codegen part of the process.

From: Mehdi Amini <mehdi.amini@apple.com>
llvm-svn: 265113

5 files changed, 83 insertions, 5 deletions

diff --git a/llvm/include/llvm-c/lto.h b/llvm/include/llvm-c/lto.h
index 42c280210a2..82f7fd74178 100644
--- a/llvm/include/llvm-c/lto.h
+++ b/llvm/include/llvm-c/lto.h
@@ -44,7 +44,8 @@ typedef bool lto_bool_t;
  * @{
  */
 
-#define LTO_API_VERSION 18
+#define LTO_API_VERSION 19
+
 /**
  * \since prior to LTO_API_VERSION=3
  */
@@ -718,6 +719,23 @@ extern void thinlto_codegen_set_savetemps_dir(thinlto_code_gen_t cg,
 extern void thinlto_codegen_set_cpu(thinlto_code_gen_t cg, const char *cpu);
 
 /**
+ * Disable CodeGen, only run the stages till codegen and stop. The output will
+ * be bitcode.
+ *
+ * \since LTO_API_VERSION=19
+ */
+extern void thinlto_codegen_disable_codegen(thinlto_code_gen_t cg,
+                                            lto_bool_t disable);
+
+/**
+ * Perform CodeGen only: disable all other stages.
+ *
+ * \since LTO_API_VERSION=19
+ */
+extern void thinlto_codegen_set_codegen_only(thinlto_code_gen_t cg,
+                                             lto_bool_t codegen_only);
+
+/**
  * Parse -mllvm style debug options.
  *
  * \since LTO_API_VERSION=18
diff --git a/llvm/include/llvm/LTO/ThinLTOCodeGenerator.h b/llvm/include/llvm/LTO/ThinLTOCodeGenerator.h
index 76dad3d8425..3407c82d3aa 100644
--- a/llvm/include/llvm/LTO/ThinLTOCodeGenerator.h
+++ b/llvm/include/llvm/LTO/ThinLTOCodeGenerator.h
@@ -169,6 +169,13 @@ public:
     TMBuilder.CGOptLevel = CGOptLevel;
   }
 
+  /// Disable CodeGen, only run the stages till codegen and stop. The output
+  /// will be bitcode.
+  void disableCodeGen(bool Disable) { DisableCodeGen = Disable; }
+
+  /// Perform CodeGen only: disable all other stages.
+  void setCodeGenOnly(bool CGOnly) { CodeGenOnly = CGOnly; }
+
   /**@}*/
 
   /**
@@ -228,6 +235,14 @@ private:
 
   /// Path to a directory to save the temporary bitcode files.
   std::string SaveTempsDir;
+
+  /// Flag to enable/disable CodeGen. When set to true, the process stops after
+  /// optimizations and a bitcode is produced.
+  bool DisableCodeGen;
+
+  /// Flag to indicate that only the CodeGen will be performed, no cross-module
+  /// importing or optimization.
+  bool CodeGenOnly;
 };
 }
 #endif
diff --git a/llvm/lib/LTO/ThinLTOCodeGenerator.cpp b/llvm/lib/LTO/ThinLTOCodeGenerator.cpp
index c80e7bdffb0..412475d1c98 100644
--- a/llvm/lib/LTO/ThinLTOCodeGenerator.cpp
+++ b/llvm/lib/LTO/ThinLTOCodeGenerator.cpp
@@ -188,7 +188,8 @@ ProcessThinLTOModule(Module &TheModule, const ModuleSummaryIndex &Index,
                      StringMap<MemoryBufferRef> &ModuleMap, TargetMachine &TM,
                      const FunctionImporter::ImportMapTy &ImportList,
                      ThinLTOCodeGenerator::CachingOptions CacheOptions,
-                     StringRef SaveTempsDir, unsigned count) {
+                     bool DisableCodeGen, StringRef SaveTempsDir,
+                     unsigned count) {
 
   // Save temps: after IPO.
   saveTempBitcode(TheModule, SaveTempsDir, count, ".1.IPO.bc");
@@ -212,6 +213,16 @@ ProcessThinLTOModule(Module &TheModule, const ModuleSummaryIndex &Index,
 
   saveTempBitcode(TheModule, SaveTempsDir, count, ".3.opt.bc");
 
+  if (DisableCodeGen) {
+    // Configured to stop before CodeGen, serialize the bitcode and return.
+    SmallVector<char, 128> OutputBuffer;
+    {
+      raw_svector_ostream OS(OutputBuffer);
+      WriteBitcodeToFile(&TheModule, OS, true, true);
+    }
+    return make_unique<ObjectMemoryBuffer>(std::move(OutputBuffer));
+  }
+
   return codegenModule(TheModule, TM);
 }
 
@@ -348,6 +359,28 @@ std::unique_ptr<MemoryBuffer> ThinLTOCodeGenerator::codegen(Module &TheModule) {
 
 // Main entry point for the ThinLTO processing
 void ThinLTOCodeGenerator::run() {
+  if (CodeGenOnly) {
+    // Perform only parallel codegen and return.
+    ThreadPool Pool;
+    assert(ProducedBinaries.empty() && "The generator should not be reused");
+    ProducedBinaries.resize(Modules.size());
+    int count = 0;
+    for (auto &ModuleBuffer : Modules) {
+      Pool.async([&](int count) {
+        LLVMContext Context;
+        Context.setDiscardValueNames(LTODiscardValueNames);
+
+        // Parse module now
+        auto TheModule = loadModuleFromBuffer(ModuleBuffer, Context, false);
+
+        // CodeGen
+        ProducedBinaries[count] = codegen(*TheModule);
+      }, count++);
+    }
+
+    return;
+  }
+
   // Sequential linking phase
   auto Index = linkCombinedIndex();
 
@@ -396,7 +429,7 @@ void ThinLTOCodeGenerator::run() {
         auto &ImportList = ImportLists[TheModule->getModuleIdentifier()];
         ProducedBinaries[count] = ProcessThinLTOModule(
             *TheModule, *Index, ModuleMap, *TMBuilder.create(), ImportList,
-            CacheOptions, SaveTempsDir, count);
+            CacheOptions, DisableCodeGen, SaveTempsDir, count);
       }, count);
       count++;
     }
diff --git a/llvm/tools/lto/lto.cpp b/llvm/tools/lto/lto.cpp
index 165b23da024..e22b198b3e1 100644
--- a/llvm/tools/lto/lto.cpp
+++ b/llvm/tools/lto/lto.cpp
@@ -473,6 +473,16 @@ LTOObjectBuffer thinlto_module_get_object(thinlto_code_gen_t cg,
                          MemBuffer->getBufferSize()};
 }
 
+void thinlto_codegen_disable_codegen(thinlto_code_gen_t cg,
+                                     lto_bool_t disable) {
+  unwrap(cg)->disableCodeGen(disable);
+}
+
+void thinlto_codegen_set_codegen_only(thinlto_code_gen_t cg,
+                                      lto_bool_t CodeGenOnly) {
+  unwrap(cg)->setCodeGenOnly(CodeGenOnly);
+}
+
 void thinlto_debug_options(const char *const *options, int number) {
   // if options were requested, set them
   if (number && options) {
@@ -483,7 +493,7 @@ void thinlto_debug_options(const char *const *options, int number) {
   }
 }
 
-bool lto_module_is_thinlto(lto_module_t mod) {
+lto_bool_t lto_module_is_thinlto(lto_module_t mod) {
   return unwrap(mod)->isThinLTO();
 }
 
diff --git a/llvm/tools/lto/lto.exports b/llvm/tools/lto/lto.exports
index e0df0ca0e4f..83a59437e81 100644
--- a/llvm/tools/lto/lto.exports
+++ b/llvm/tools/lto/lto.exports
@@ -61,4 +61,6 @@ thinlto_debug_options
 lto_module_is_thinlto
 thinlto_codegen_add_must_preserve_symbol
 thinlto_codegen_add_cross_referenced_symbol
-thinlto_codegen_set_final_cache_size_relative_to_available_space
\ No newline at end of file
+thinlto_codegen_set_final_cache_size_relative_to_available_space
+thinlto_codegen_set_codegen_only
+thinlto_codegen_disable_codegen
\ No newline at end of file