[lld-link] implement -thinlto-{prefix,object-suffix}-replace

Summary:
Adds the following two options to lld-link:

-thinlto-prefix-replace: allows replacing a prefix in paths generated
for ThinLTO. This can be used to ensure index files and native object
files are stored in unique directories, allowing multiple distributed
ThinLTO links to proceed concurrently.

-thinlto-object-suffix-replace: allows replacing a suffix in object
file paths involved in ThinLTO. This allows minimized index files to
be used for the thin link while storing the paths to the full bitcode
files for subsequent steps (code generation and final linking).

Reviewers: ruiu, tejohnson, pcc, rnk

Subscribers: mehdi_amini, steven_wu, dexonsmith, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D64542

llvm-svn: 365807

8 files changed, 131 insertions, 5 deletions

diff --git a/lld/COFF/Config.h b/lld/COFF/Config.h
index 60684726c06..6d9d6a848da 100644
--- a/lld/COFF/Config.h
+++ b/lld/COFF/Config.h
@@ -182,6 +182,12 @@ struct Configuration {
   // Used for /thinlto-index-only:
   llvm::StringRef thinLTOIndexOnlyArg;
 
+  // Used for /thinlto-object-prefix-replace:
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOPrefixReplace;
+
+  // Used for /thinlto-object-suffix-replace:
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOObjectSuffixReplace;
+
   uint64_t imageBase = -1;
   uint64_t fileAlign = 512;
   uint64_t stackReserve = 1024 * 1024;
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 304fc1837b7..0009296124e 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -86,6 +86,20 @@ bool link(ArrayRef<const char *> args, bool canExitEarly, raw_ostream &diag) {
   return !errorCount();
 }
 
+// Parse options of the form "old;new".
+static std::pair<StringRef, StringRef> getOldNewOptions(opt::InputArgList &args,
+                                                        unsigned id) {
+  auto *arg = args.getLastArg(id);
+  if (!arg)
+    return {"", ""};
+
+  StringRef s = arg->getValue();
+  std::pair<StringRef, StringRef> ret = s.split(';');
+  if (ret.second.empty())
+    error(arg->getSpelling() + " expects 'old;new' format, but got " + s);
+  return ret;
+}
+
 // Drop directory components and replace extension with ".exe" or ".dll".
 static std::string getOutputPath(StringRef path) {
   auto p = path.find_last_of("\\/");
@@ -1446,6 +1460,10 @@ void LinkerDriver::link(ArrayRef<const char *> argsArr) {
                              args.hasArg(OPT_thinlto_index_only_arg);
   config->thinLTOIndexOnlyArg =
       args.getLastArgValue(OPT_thinlto_index_only_arg);
+  config->thinLTOPrefixReplace =
+      getOldNewOptions(args, OPT_thinlto_prefix_replace);
+  config->thinLTOObjectSuffixReplace =
+      getOldNewOptions(args, OPT_thinlto_object_suffix_replace);
   // Handle miscellaneous boolean flags.
   config->allowBind = args.hasFlag(OPT_allowbind, OPT_allowbind_no, true);
   config->allowIsolation =
diff --git a/lld/COFF/InputFiles.cpp b/lld/COFF/InputFiles.cpp
index 6737db884ed..e28d07a718b 100644
--- a/lld/COFF/InputFiles.cpp
+++ b/lld/COFF/InputFiles.cpp
@@ -783,6 +783,8 @@ BitcodeFile::BitcodeFile(MemoryBufferRef mb, StringRef archiveName,
                          uint64_t offsetInArchive)
     : InputFile(BitcodeKind, mb) {
   std::string path = mb.getBufferIdentifier().str();
+  if (config->thinLTOIndexOnly)
+    path = replaceThinLTOSuffix(mb.getBufferIdentifier());
 
   // ThinLTO assumes that all MemoryBufferRefs given to it have a unique
   // name. If two archives define two members with the same name, this
@@ -849,6 +851,15 @@ MachineTypes BitcodeFile::getMachineType() {
     return IMAGE_FILE_MACHINE_UNKNOWN;
   }
 }
+
+std::string replaceThinLTOSuffix(StringRef path) {
+  StringRef suffix = config->thinLTOObjectSuffixReplace.first;
+  StringRef repl = config->thinLTOObjectSuffixReplace.second;
+
+  if (path.consume_back(suffix))
+    return (path + repl).str();
+  return path;
+}
 } // namespace coff
 } // namespace lld
 
diff --git a/lld/COFF/InputFiles.h b/lld/COFF/InputFiles.h
index bd2698f4761..6c84ceec13d 100644
--- a/lld/COFF/InputFiles.h
+++ b/lld/COFF/InputFiles.h
@@ -311,6 +311,8 @@ private:
 
   std::vector<Symbol *> symbols;
 };
+
+std::string replaceThinLTOSuffix(StringRef path);
 } // namespace coff
 
 std::string toString(const coff::InputFile *file);
diff --git a/lld/COFF/LTO.cpp b/lld/COFF/LTO.cpp
index 3ebe0b881ec..eb3c60d6607 100644
--- a/lld/COFF/LTO.cpp
+++ b/lld/COFF/LTO.cpp
@@ -54,6 +54,12 @@ static std::unique_ptr<raw_fd_ostream> openFile(StringRef file) {
   return ret;
 }
 
+static std::string getThinLTOOutputFile(StringRef path) {
+  return lto::getThinLTOOutputFile(path,
+                                   config->thinLTOPrefixReplace.first,
+                                   config->thinLTOPrefixReplace.second);
+}
+
 static lto::Config createConfig() {
   lto::Config c;
   c.Options = initTargetOptionsFromCodeGenFlags();
@@ -93,7 +99,8 @@ BitcodeCompiler::BitcodeCompiler() {
   if (config->thinLTOIndexOnly) {
     auto OnIndexWrite = [&](StringRef S) { thinIndices.erase(S); };
     backend = lto::createWriteIndexesThinBackend(
-        "", "", config->thinLTOEmitImportsFiles, indexFile.get(), OnIndexWrite);
+        config->thinLTOPrefixReplace.first, config->thinLTOPrefixReplace.second,
+        config->thinLTOEmitImportsFiles, indexFile.get(), OnIndexWrite);
   } else if (config->thinLTOJobs != 0) {
     backend = lto::createInProcessThinBackend(config->thinLTOJobs);
   }
@@ -159,11 +166,11 @@ std::vector<StringRef> BitcodeCompiler::compile() {
       cache));
 
   // Emit empty index files for non-indexed files
-  for (StringRef S : thinIndices) {
-    std::string Path(S);
-    openFile(Path + ".thinlto.bc");
+  for (StringRef s : thinIndices) {
+    std::string path = getThinLTOOutputFile(s);
+    openFile(path + ".thinlto.bc");
     if (config->thinLTOEmitImportsFiles)
-      openFile(Path + ".imports");
+      openFile(path + ".imports");
   }
 
   // ThinLTO with index only option is required to generate only the index
diff --git a/lld/COFF/Options.td b/lld/COFF/Options.td
index 5f54903e88e..024b7be8f78 100644
--- a/lld/COFF/Options.td
+++ b/lld/COFF/Options.td
@@ -185,6 +185,12 @@ def thinlto_index_only :
 def thinlto_index_only_arg : P<
     "thinlto-index-only",
     "-thinlto-index-only and also write native module names to file">;
+def thinlto_object_suffix_replace : P<
+    "thinlto-object-suffix-replace",
+    "'old;new' replace old suffix with new suffix in ThinLTO index">;
+def thinlto_prefix_replace: P<
+    "thinlto-prefix-replace",
+    "'old;new' replace old prefix with new prefix in ThinLTO outputs">;
 def dash_dash_version : Flag<["--"], "version">,
   HelpText<"Print version information">;
 defm threads: B<"threads",
diff --git a/lld/test/COFF/thinlto-object-suffix-replace.ll b/lld/test/COFF/thinlto-object-suffix-replace.ll
new file mode 100644
index 00000000000..eea898de0c5
--- /dev/null
+++ b/lld/test/COFF/thinlto-object-suffix-replace.ll
@@ -0,0 +1,50 @@
+; REQUIRES: x86
+
+; Test to make sure the thinlto-object-suffix-replace option is handled
+; correctly.
+
+; Generate bitcode file with summary, as well as a minimized bitcode without
+; the debug metadata for the thin link.
+; RUN: opt -thinlto-bc %s -thin-link-bitcode-file=%t1.thinlink.bc -o %t1.obj
+
+; First perform the thin link on the normal bitcode file, and save the
+; resulting index.
+; RUN: lld-link -thinlto-index-only -entry:main %t1.obj -out:%t3.exe
+; RUN: cp %t1.obj.thinlto.bc %t1.obj.thinlto.bc.orig
+
+; Next perform the thin link on the minimized bitcode file, and compare dump
+; of the resulting index to the above dump to ensure they are identical.
+; RUN: rm -f %t1.obj.thinlto.bc
+; Make sure it isn't inadvertently using the regular bitcode file.
+; RUN: rm -f %t1.obj
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-object-suffix-replace:".thinlink.bc;.obj" \
+; RUN:     %t1.thinlink.bc -out:%t3.exe
+; RUN: diff %t1.obj.thinlto.bc.orig %t1.obj.thinlto.bc
+
+; Ensure lld generates error if suffix replace option not in 'old;new' format.
+; RUN: rm -f %t1.obj.thinlto.bc
+; RUN: not lld-link -entry:main -thinlto-index-only \
+; RUN: -thinlto-object-suffix-replace:"abc:def" %t1.thinlink.bc \
+; RUN: -out:%t3.exe 2>&1 | FileCheck %s --check-prefix=ERR1
+; ERR1: -thinlto-object-suffix-replace: expects 'old;new' format, but got abc:def
+
+; If filename does not end with old suffix, no suffix change should occur,
+; so ".thinlto.bc" will simply be appended to the input file name.
+; RUN: rm -f %t1.thinlink.bc.thinlto.bc
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN: -thinlto-object-suffix-replace:".abc;.obj" %t1.thinlink.bc -out:%t3.exe
+; RUN: ls %t1.thinlink.bc.thinlto.bc
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.24215"
+
+define void @main() {
+entry:
+  ret void
+}
+
+!llvm.dbg.cu = !{}
+
+!1 = !{i32 2, !"Debug Info Version", i32 3}
+!llvm.module.flags = !{!1}
diff --git a/lld/test/COFF/thinlto-prefix-replace.ll b/lld/test/COFF/thinlto-prefix-replace.ll
new file mode 100644
index 00000000000..a6a116d67e1
--- /dev/null
+++ b/lld/test/COFF/thinlto-prefix-replace.ll
@@ -0,0 +1,26 @@
+; REQUIRES: x86
+; Check that changing the output path via thinlto-prefix-replace works
+; RUN: mkdir -p %t/oldpath
+; RUN: opt -module-summary %s -o %t/oldpath/t.obj
+
+; Ensure that there is no existing file at the new path, so we properly
+; test the creation of the new file there.
+; RUN: rm -f %t/newpath/t.obj.thinlto.bc
+; RUN: lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-prefix-replace:"%t/oldpath/;%t/newpath/" %t/oldpath/t.obj \
+; RUN:     -out:%t/t.exe
+; RUN: ls %t/newpath/t.obj.thinlto.bc
+
+; Ensure that lld errors if prefix replace option is not in 'old;new' format.
+; RUN: rm -f %t/newpath/t.obj.thinlto.bc
+; RUN: not lld-link -entry:main -thinlto-index-only \
+; RUN:     -thinlto-prefix-replace:"abc:def" %t/oldpath/t.obj \
+; RUN:     -out:%t/t.exe 2>&1 | FileCheck --check-prefix=ERR %s
+; ERR: -thinlto-prefix-replace: expects 'old;new' format, but got abc:def
+
+target datalayout = "e-m:w-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-windows-msvc19.0.24215"
+
+define void @main() {
+  ret void
+}