[ThinLTO] Option to control path of distributed backend files

Summary:
Add support to control where files for a distributed backend (the
individual index files and optional imports files) are created.

This is invoked with a new thinlto-prefix-replace option in the gold
plugin and llvm-lto. If specified, expects a string of the form
"oldprefix:newprefix", and instead of generating these files in the
same directory path as the corresponding bitcode file, will use a path
formed by replacing the bitcode file's path prefix matching oldprefix
with newprefix.

Also add a new replace_path_prefix helper to Path.h in libSupport.

Depends on D19636.

Reviewers: joker.eph

Subscribers: llvm-commits, joker.eph

Differential Revision: http://reviews.llvm.org/D19644

llvm-svn: 269771

llvm/include/llvm/Support/Path.h

@@ -142,6 +142,23 @@ void remove_filename(SmallVectorImpl<char> &path);
 ///                  prepended.
 void replace_extension(SmallVectorImpl<char> &path, const Twine &extension);
 
+/// @brief Replace matching path prefix with another path.
+///
+/// @code
+///   /foo, /old, /new => /foo
+///   /old/foo, /old, /new => /new/foo
+///   /foo, <empty>, /new => /new/foo
+///   /old/foo, /old, <empty> => /foo
+/// @endcode
+///
+/// @param Path If \a Path starts with \a OldPrefix modify to instead
+///        start with \a NewPrefix.
+/// @param OldPrefix The path prefix to strip from \a Path.
+/// @param NewPrefix The path prefix to replace \a NewPrefix with.
+void replace_path_prefix(SmallVectorImpl<char> &Path,
+                         const StringRef &OldPrefix,
+                         const StringRef &NewPrefix);
+
 /// @brief Append to path.
 ///
 /// @code

llvm/lib/Support/Path.cpp

@@ -522,6 +522,29 @@ void replace_extension(SmallVectorImpl<char> &path, const Twine &extension) {
   path.append(ext.begin(), ext.end());
 }
 
+void replace_path_prefix(SmallVectorImpl<char> &Path,
+                         const StringRef &OldPrefix,
+                         const StringRef &NewPrefix) {
+  if (OldPrefix.empty() && NewPrefix.empty())
+    return;
+
+  StringRef OrigPath(Path.begin(), Path.size());
+  if (!OrigPath.startswith(OldPrefix))
+    return;
+
+  // If prefixes have the same size we can simply copy the new one over.
+  if (OldPrefix.size() == NewPrefix.size()) {
+    std::copy(NewPrefix.begin(), NewPrefix.end(), Path.begin());
+    return;
+  }
+
+  StringRef RelPath = OrigPath.substr(OldPrefix.size());
+  SmallString<256> NewPath;
+  path::append(NewPath, NewPrefix);
+  path::append(NewPath, RelPath);
+  Path.swap(NewPath);
+}
+
 void native(const Twine &path, SmallVectorImpl<char> &result) {
   assert((!path.isSingleStringRef() ||
           path.getSingleStringRef().data() != result.data()) &&

llvm/test/ThinLTO/X86/prefix_replace.ll

@@ -0,0 +1,16 @@
+; Check that changing the output path via prefix-replace works
+; RUN: mkdir -p %T/oldpath
+; RUN: opt -module-summary %s -o %T/oldpath/prefix_replace.o
+; 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/prefix_replace.o.thinlto.bc
+
+; RUN: llvm-lto -thinlto-action=thinlink -o %t.index.bc %T/oldpath/prefix_replace.o
+; RUN: llvm-lto -thinlto-action=distributedindexes -thinlto-prefix-replace="%T/oldpath/:%T/newpath/" -thinlto-index %t.index.bc %T/oldpath/prefix_replace.o
+
+; RUN: ls %T/newpath/prefix_replace.o.thinlto.bc
+
+define void @f() {
+entry:
+  ret void
+}

llvm/test/tools/gold/X86/thinlto_prefix_replace.ll

@@ -0,0 +1,17 @@
+; Check that changing the output path via thinlto-prefix-replace works
+; RUN: mkdir -p %T/oldpath
+; RUN: opt -module-summary %s -o %T/oldpath/thinlto_prefix_replace.o
+; 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/thinlto_prefix_replace.o.thinlto.bc
+; RUN: %gold -plugin %llvmshlibdir/LLVMgold.so \
+; RUN:    --plugin-opt=thinlto \
+; RUN:    --plugin-opt=thinlto-index-only \
+; RUN:    --plugin-opt=thinlto-prefix-replace="%T/oldpath/:%T/newpath/" \
+; RUN:    -shared %T/oldpath/thinlto_prefix_replace.o -o %T/thinlto_prefix_replace
+; RUN: ls %T/newpath/thinlto_prefix_replace.o.thinlto.bc
+
+define void @f() {
+entry:
+  ret void
+}

llvm/tools/gold/gold-plugin.cpp

@@ -35,6 +35,7 @@
 #include "llvm/Support/Host.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/ThreadPool.h"
@@ -189,6 +190,13 @@ namespace options {
   // bitcode file, listing the files it imports from in plain text. This is to
   // support distributed build file staging.
   static bool thinlto_emit_imports_files = false;
+  // Option to control where files for a distributed backend (the individual
+  // index files and optional imports files) are created.
+  // If specified, expects a string of the form "oldprefix:newprefix", and
+  // instead of generating these files in the same directory path as the
+  // corresponding bitcode file, will use a path formed by replacing the
+  // bitcode file's path prefix matching oldprefix with newprefix.
+  static std::string thinlto_prefix_replace;
   // Additional options to pass into the code generator.
   // Note: This array will contain all plugin options which are not claimed
   // as plugin exclusive to pass to the code generator.
@@ -224,6 +232,10 @@ namespace options {
       thinlto_index_only = true;
     } else if (opt == "thinlto-emit-imports-files") {
       thinlto_emit_imports_files = true;
+    } else if (opt.startswith("thinlto-prefix-replace=")) {
+      thinlto_prefix_replace = opt.substr(strlen("thinlto-prefix-replace="));
+      if (thinlto_prefix_replace.find(":") == std::string::npos)
+        message(LDPL_FATAL, "thinlto-prefix-replace expects 'old:new' format");
     } else if (opt.size() == 2 && opt[0] == 'O') {
       if (opt[1] < '0' || opt[1] > '3')
         message(LDPL_FATAL, "Optimization level must be between 0 and 3");
@@ -1202,6 +1214,37 @@ static void thinLTOBackends(raw_fd_ostream *ApiFile,
     Task.cleanup();
 }
 
+/// Parse the thinlto_prefix_replace option into the \p OldPrefix and
+/// \p NewPrefix strings, if it was specified.
+static void getThinLTOOldAndNewPrefix(std::string &OldPrefix,
+                                      std::string &NewPrefix) {
+  StringRef PrefixReplace = options::thinlto_prefix_replace;
+  assert(PrefixReplace.empty() || PrefixReplace.find(":") != StringRef::npos);
+  std::pair<StringRef, StringRef> Split = PrefixReplace.split(":");
+  OldPrefix = Split.first.str();
+  NewPrefix = Split.second.str();
+}
+
+/// Given the original \p Path to an output file, replace any path
+/// prefix matching \p OldPrefix with \p NewPrefix. Also, create the
+/// resulting directory if it does not yet exist.
+static std::string getThinLTOOutputFile(const std::string &Path,
+                                        const std::string &OldPrefix,
+                                        const std::string &NewPrefix) {
+  if (OldPrefix.empty() && NewPrefix.empty())
+    return Path;
+  SmallString<128> NewPath(Path);
+  llvm::sys::path::replace_path_prefix(NewPath, OldPrefix, NewPrefix);
+  StringRef ParentPath = llvm::sys::path::parent_path(NewPath.str());
+  if (!ParentPath.empty()) {
+    // Make sure the new directory exists, creating it if necessary.
+    if (std::error_code EC = llvm::sys::fs::create_directories(ParentPath))
+      llvm::errs() << "warning: could not create directory '" << ParentPath
+                   << "': " << EC.message() << '\n';
+  }
+  return NewPath.str();
+}
+
 /// Perform ThinLTO link, which creates the combined index file.
 /// Also, either launch backend threads or (under thinlto-index-only)
 /// emit individual index files for distributed backends and exit.
@@ -1240,17 +1283,25 @@ static ld_plugin_status thinLTOLink(raw_fd_ostream *ApiFile) {
     ComputeCrossModuleImport(CombinedIndex, ModuleToDefinedGVSummaries,
                              ImportLists, ExportLists);
 
+    // If the thinlto-prefix-replace option was specified, parse it and
+    // extract the old and new prefixes.
+    std::string OldPrefix, NewPrefix;
+    getThinLTOOldAndNewPrefix(OldPrefix, NewPrefix);
+
     // For each input bitcode file, generate an individual index that
     // contains summaries only for its own global values, and for any that
     // should be imported.
     for (claimed_file &F : Modules) {
       PluginInputFile InputFile(F.handle);
       std::error_code EC;
-      raw_fd_ostream OS((Twine(InputFile.file().name) + ".thinlto.bc").str(),
-                        EC, sys::fs::OpenFlags::F_None);
+
+      std::string NewModulePath =
+          getThinLTOOutputFile(InputFile.file().name, OldPrefix, NewPrefix);
+      raw_fd_ostream OS((Twine(NewModulePath) + ".thinlto.bc").str(), EC,
+                        sys::fs::OpenFlags::F_None);
       if (EC)
         message(LDPL_FATAL, "Unable to open %s.thinlto.bc for writing: %s",
-                InputFile.file().name, EC.message().c_str());
+                NewModulePath.c_str(), EC.message().c_str());
       // Build a map of module to the GUIDs and summary objects that should
       // be written to its index.
       std::map<std::string, GVSummaryMapTy> ModuleToSummariesForIndex;
@@ -1262,10 +1313,10 @@ static ld_plugin_status thinLTOLink(raw_fd_ostream *ApiFile) {
       if (options::thinlto_emit_imports_files) {
         if ((EC = EmitImportsFiles(
                  InputFile.file().name,
-                 (Twine(InputFile.file().name) + ".imports").str(),
+                 (Twine(NewModulePath) + ".imports").str(),
                  ImportLists)))
           message(LDPL_FATAL, "Unable to open %s.imports",
-                  InputFile.file().name, EC.message().c_str());
+                  NewModulePath.c_str(), EC.message().c_str());
       }
     }
 

llvm/tools/llvm-lto/llvm-lto.cpp

@@ -26,6 +26,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/SourceMgr.h"
@@ -104,6 +105,13 @@ static cl::opt<std::string>
                  cl::desc("Provide the index produced by a ThinLink, required "
                           "to perform the promotion and/or importing."));
 
+static cl::opt<std::string> ThinLTOPrefixReplace(
+    "thinlto-prefix-replace",
+    cl::desc("Control where files for distributed backends are "
+             "created. Expects 'oldprefix:newprefix' and if path "
+             "prefix of output file is oldprefix it will be "
+             "replaced with newprefix."));
+
 static cl::opt<std::string> ThinLTOModuleId(
     "thinlto-module-id",
     cl::desc("For the module ID for the file to process, useful to "
@@ -294,6 +302,37 @@ static void createCombinedModuleSummaryIndex() {
   OS.close();
 }
 
+/// Parse the thinlto_prefix_replace option into the \p OldPrefix and
+/// \p NewPrefix strings, if it was specified.
+static void getThinLTOOldAndNewPrefix(std::string &OldPrefix,
+                                      std::string &NewPrefix) {
+  assert(ThinLTOPrefixReplace.empty() ||
+         ThinLTOPrefixReplace.find(":") != StringRef::npos);
+  StringRef PrefixReplace = ThinLTOPrefixReplace;
+  std::pair<StringRef, StringRef> Split = PrefixReplace.split(":");
+  OldPrefix = Split.first.str();
+  NewPrefix = Split.second.str();
+}
+
+/// Given the original \p Path to an output file, replace any path
+/// prefix matching \p OldPrefix with \p NewPrefix. Also, create the
+/// resulting directory if it does not yet exist.
+static std::string getThinLTOOutputFile(const std::string &Path,
+                                        const std::string &OldPrefix,
+                                        const std::string &NewPrefix) {
+  if (OldPrefix.empty() && NewPrefix.empty())
+    return Path;
+  SmallString<128> NewPath(Path);
+  llvm::sys::path::replace_path_prefix(NewPath, OldPrefix, NewPrefix);
+  StringRef ParentPath = llvm::sys::path::parent_path(NewPath.str());
+  if (!ParentPath.empty()) {
+    // Make sure the new directory exists, creating it if necessary.
+    if (std::error_code EC = llvm::sys::fs::create_directories(ParentPath))
+      error(EC, "error creating the directory '" + ParentPath + "'");
+  }
+  return NewPath.str();
+}
+
 namespace thinlto {
 
 std::vector<std::unique_ptr<MemoryBuffer>>
@@ -421,6 +460,9 @@ private:
                          "the output files will be suffixed from the input "
                          "ones.");
 
+    std::string OldPrefix, NewPrefix;
+    getThinLTOOldAndNewPrefix(OldPrefix, NewPrefix);
+
     auto Index = loadCombinedIndex();
     for (auto &Filename : InputFilenames) {
       // Build a map of module to the GUIDs and summary objects that should
@@ -433,6 +475,7 @@ private:
       if (OutputName.empty()) {
         OutputName = Filename + ".thinlto.bc";
       }
+      OutputName = getThinLTOOutputFile(OutputName, OldPrefix, NewPrefix);
       std::error_code EC;
       raw_fd_ostream OS(OutputName, EC, sys::fs::OpenFlags::F_None);
       error(EC, "error opening the file '" + OutputName + "'");
@@ -449,12 +492,16 @@ private:
                          "the output files will be suffixed from the input "
                          "ones.");
 
+    std::string OldPrefix, NewPrefix;
+    getThinLTOOldAndNewPrefix(OldPrefix, NewPrefix);
+
     auto Index = loadCombinedIndex();
     for (auto &Filename : InputFilenames) {
       std::string OutputName = OutputFilename;
       if (OutputName.empty()) {
         OutputName = Filename + ".imports";
       }
+      OutputName = getThinLTOOutputFile(OutputName, OldPrefix, NewPrefix);
       ThinLTOCodeGenerator::emitImports(Filename, OutputName, *Index);
     }
   }

llvm/unittests/Support/Path.cpp

@@ -970,4 +970,29 @@ TEST(Support, RemoveDots) {
   EXPECT_EQ("c", Path1);
 #endif
 }
+
+TEST(Support, ReplacePathPrefix) {
+  SmallString<64> Path1("/foo");
+  SmallString<64> Path2("/old/foo");
+  SmallString<64> OldPrefix("/old");
+  SmallString<64> NewPrefix("/new");
+  SmallString<64> NewPrefix2("/longernew");
+  SmallString<64> EmptyPrefix("");
+
+  SmallString<64> Path = Path1;
+  path::replace_path_prefix(Path, OldPrefix, NewPrefix);
+  EXPECT_EQ(Path, "/foo");
+  Path = Path2;
+  path::replace_path_prefix(Path, OldPrefix, NewPrefix);
+  EXPECT_EQ(Path, "/new/foo");
+  Path = Path2;
+  path::replace_path_prefix(Path, OldPrefix, NewPrefix2);
+  EXPECT_EQ(Path, "/longernew/foo");
+  Path = Path1;
+  path::replace_path_prefix(Path, EmptyPrefix, NewPrefix);
+  EXPECT_EQ(Path, "/new/foo");
+  Path = Path2;
+  path::replace_path_prefix(Path, OldPrefix, EmptyPrefix);
+  EXPECT_EQ(Path, "/foo");
+}
 } // anonymous namespace