1 files changed, 45 insertions, 4 deletions

diff --git a/llvm/tools/llvm-lto/llvm-lto.cpp b/llvm/tools/llvm-lto/llvm-lto.cpp
index 98644d3078f..3b13c7d1640 100644
--- a/llvm/tools/llvm-lto/llvm-lto.cpp
+++ b/llvm/tools/llvm-lto/llvm-lto.cpp
@@ -68,6 +68,7 @@ enum ThinLTOModes {
   THINLINK,
   THINPROMOTE,
   THINIMPORT,
+  THININTERNALIZE,
   THINOPT,
   THINCODEGEN,
   THINALL
@@ -84,6 +85,9 @@ cl::opt<ThinLTOModes> ThinLTOMode(
         clEnumValN(THINIMPORT, "import", "Perform both promotion and "
                                          "cross-module importing (requires "
                                          "-thinlto-index)."),
+        clEnumValN(THININTERNALIZE, "internalize",
+                   "Perform internalization driven by -exported-symbol "
+                   "(requires -thinlto-index)."),
         clEnumValN(THINOPT, "optimize", "Perform ThinLTO optimizations."),
         clEnumValN(THINCODEGEN, "codegen", "CodeGen (expected to match llc)"),
         clEnumValN(THINALL, "run", "Perform ThinLTO end-to-end"),
@@ -105,10 +109,10 @@ static cl::opt<std::string> OutputFilename("o", cl::init(""),
                                            cl::desc("Override output filename"),
                                            cl::value_desc("filename"));
 
-static cl::list<std::string>
-    ExportedSymbols("exported-symbol",
-                    cl::desc("Symbol to export from the resulting object file"),
-                    cl::ZeroOrMore);
+static cl::list<std::string> ExportedSymbols(
+    "exported-symbol",
+    cl::desc("List of symbols to export from the resulting object file"),
+    cl::ZeroOrMore);
 
 static cl::list<std::string>
     DSOSymbols("dso-symbol",
@@ -329,6 +333,10 @@ public:
   ThinLTOProcessing(const TargetOptions &Options) {
     ThinGenerator.setCodePICModel(RelocModel);
     ThinGenerator.setTargetOptions(Options);
+
+    // Add all the exported symbols to the table of symbols to preserve.
+    for (unsigned i = 0; i < ExportedSymbols.size(); ++i)
+      ThinGenerator.preserveSymbol(ExportedSymbols[i]);
   }
 
   void run() {
@@ -339,6 +347,8 @@ public:
       return promote();
     case THINIMPORT:
       return import();
+    case THININTERNALIZE:
+      return internalize();
     case THINOPT:
       return optimize();
     case THINCODEGEN:
@@ -432,6 +442,37 @@ private:
     }
   }
 
+  void internalize() {
+    if (InputFilenames.size() != 1 && !OutputFilename.empty())
+      report_fatal_error("Can't handle a single output filename and multiple "
+                         "input files, do not provide an output filename and "
+                         "the output files will be suffixed from the input "
+                         "ones.");
+
+    if (ExportedSymbols.empty())
+      errs() << "Warning: -internalize will not perform without "
+                "-exported-symbol\n";
+
+    auto Index = loadCombinedIndex();
+    auto InputBuffers = loadAllFilesForIndex(*Index);
+    for (auto &MemBuffer : InputBuffers)
+      ThinGenerator.addModule(MemBuffer->getBufferIdentifier(),
+                              MemBuffer->getBuffer());
+
+    for (auto &Filename : InputFilenames) {
+      LLVMContext Ctx;
+      auto TheModule = loadModule(Filename, Ctx);
+
+      ThinGenerator.internalize(*TheModule, *Index);
+
+      std::string OutputName = OutputFilename;
+      if (OutputName.empty()) {
+        OutputName = Filename + ".thinlto.internalized.bc";
+      }
+      writeModuleToFile(*TheModule, OutputName);
+    }
+  }
+
   void optimize() {
     if (InputFilenames.size() != 1 && !OutputFilename.empty())
       report_fatal_error("Can't handle a single output filename and multiple "