[LTO][ELF] Add --lto-exit-on option

[Northbadge]

Summary

Allows the linker to "exit early" after running specific LTO module hooks.

This is useful in combination with --save-temps= (D127778) as it gives the flexibility of stopping lld once the desired files are emitted. This minimizes wasted time (from performing optimizations, codegen, linking, etc) when all that is desired from the linker is the "--save-temps=import" output (which is completed before the aforementioned tasks), as in the case of extracting a training corpus for MLGO.

[MaskRay]

lld/ELF/Driver.cpp

1112

"unknown --lto-exit-on: " + xxx. If unknown, config->ltoExitOn should be empty.

[Northbadge]

Resolved inline comments, synced with parent, and removed -find from the test case. Thanks for the reviews!

[tejohnson]

lld/test/ELF/lto/lto-exit-on.ll

24

Why are we getting one when exiting on postpromote?

35

Why 3 here and below?

[Northbadge]

lld/test/ELF/lto/lto-exit-on.ll

24

the hook is called from lto::thinBackend AND also from LTO::runRegularLTO. When running the test case, LTO::runRegularLTO is called before lto::thinBackend from here. With postpromote being a ThinLTO exclusive hook, the postinternalize hook (from regularLTO) manages to run first before the postpromote hook is hit, thus creating 1 file. (I believe this is unavoidable without changes to that last snippet of code)

Now that you've pointed it out though, I realize that specifying --lto-exit-on=postinternalize would be ambiguous due to the two callsites, this is also the case for the preopt hook as I've just checked. The quick fix would be to make these two hooks invalid arguments for the option. Would this be an ok solution?

35

see above (that gives 1 "internalize" file)
and ThinLTO outputs the other 2, one for main.o, one for thin1.o

[Northbadge]

Removed support for --lto-exit-on={preopt,postinternalize} due to ambiguity when running ThinLTO.

[Northbadge]

lld/test/ELF/lto/lto-exit-on.ll

24

I've removed it for now, let me know if I should try doing something else

[Northbadge]

Update test to use subdirectories.

[tejohnson]

lld/test/ELF/lto/lto-exit-on.ll

24

I think it would be ok to leave the other cases in and just document it.

llvm/lib/LTO/LTOBackend.cpp

197

Probably want to issue an error here if it is something else.

[Northbadge]

Bring back deleted options, updated to be consistent with parent diff, added llvm-lto2 tests.

[Northbadge]

Missed a word, and braces

[Northbadge]

Fix test case from missing word

[MaskRay]

The summary does not mention the motivation adding --lto-exit-on. Why is useful enough to have such an option?

[Northbadge]

In D127779#3637315, @MaskRay wrote:

The summary does not mention the motivation adding --lto-exit-on. Why is useful enough to have such an option?

Added some detail on that, thanks for pointing it out

[MaskRay]

I can see that fine-grained --save-temps=xxx is useful, but the utility of --lto-exit-on is unclear. Why is the overhead considered so high that an option needs to be added?

[Northbadge]

In D127779#3637825, @MaskRay wrote:

I can see that fine-grained --save-temps=xxx is useful, but the utility of --lto-exit-on is unclear. Why is the overhead considered so high that an option needs to be added?

It is admittedly more of a nice-to-have development tool/utility. On a 128 thread machine (with enough memory) the difference between having it vs not having it for the "exit on import" case with Chrome (~30k modules) is 1m20s vs 2m20s; extrapolating the parallelized portion to a more conventional 16 thread cpu that'd be a ~5 minute difference instead

[MaskRay]

In D127779#3637851, @Northbadge wrote:

In D127779#3637825, @MaskRay wrote:

I can see that fine-grained --save-temps=xxx is useful, but the utility of --lto-exit-on is unclear. Why is the overhead considered so high that an option needs to be added?

It is admittedly more of a nice-to-have development tool/utility. On a 128 thread machine (with enough memory) the difference between having it vs not having it for the "exit on import" case with Chrome (~30k modules) is 1m20s vs 2m20s; extrapolating the parallelized portion to a more conventional 16 thread cpu that'd be a ~5 minute difference instead

And this seems specific to in-process ThinLTO. A distributed ThinLTO (usually they should use this if their executable is large) ThinLTO does not need this. So this feature seems more restricted.

[Northbadge]

An update on this: I found a different use for --lto-exit-on (well, the original intended use: in MLGO training cycles which run >1M times, which I thought wouldn't have worked initially). I'm trial-ing to make sure it's doing what I'm expecting so we can probably put this patch on hold for a week or two, it isn't super-critical at the moment. I'll report back then (+details) and see how it goes from there

[Northbadge]

It's been a bit more than 2 weeks, but the idea I had didn't exactly pan out.

For future reference, in case someone revisits this: The idea was to concatenate the postimport .bc file with its associated thinlto index file (which is also bitcode). The naive way of doing this was to use llvm-dis on both the bitcode files, append the thinlto.ll lines to the postimport.ll lines and use llvm-as to go back to bitcode. This worked on small toy examples, but crashed on larger, real, files (Something about the asm/bitcode reader/writer, iirc, didn't have time to look into it further). The motivation for doing this was so that we can run lld with a single "hermetic" module which included the necessary index entries to prevent the deletion of locally-dead(not used within own module but used in other modules) functions. This would've allowed MLGO to use lld to perfectly replicate the ThinLTO environment when it is used in a non-distributed manner, and would've necessitated this patch to prevent the linker from crashing due to missing symbols.

Abandoning the patch as I won't be working on this in the near future. Thanks for all the feedback and guidance!

lld/ELF/Config.h

struct Configuration {
   bool gnuUnique;
   bool hasDynSymTab;
   bool ignoreDataAddressEquality;
   bool ignoreFunctionAddressEquality;
   bool ltoCSProfileGenerate;
   bool ltoPGOWarnMismatch;
   bool ltoDebugPassManager;
   bool ltoEmitAsm;
+  llvm::StringRef ltoExitOn;
   bool ltoUniqueBasicBlockSectionNames;
   bool ltoWholeProgramVisibility;
   bool mergeArmExidx;
   bool mipsN32Abi = false;
   bool mmapOutputFile;
   bool nmagic;
   bool noDynamicLinker = false;
   bool noinhibitExec;

lld/ELF/Driver.cpp

for (auto *arg : args.filtered(OPT_z))
     if (!isKnownZFlag(arg->getValue()))
       warn("unknown -z value: " + StringRef(arg->getValue()));
 }
 
 constexpr const char *saveTempsValues[] = {
     "resolution", "preopt",     "promote", "internalize",  "import",
     "opt",        "precodegen", "prelink", "combinedindex"};
 
+constexpr const char *ltoExitOnValues[] = {"preopt", "promote", "internalize",
+                                           "import", "opt",     "precodegen",
+                                           "prelink"};
+
 void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   ELFOptTable parser;
   opt::InputArgList args = parser.parse(argsArr.slice(1));
 
   // Interpret these flags early because error()/warn() depend on them.
   errorHandler().errorLimit = args::getInteger(args, OPT_error_limit, 20);
   errorHandler().fatalWarnings =
       args.hasFlag(OPT_fatal_warnings, OPT_no_fatal_warnings, false);
static void readConfigs(opt::InputArgList &args) {
   config->init = args.getLastArgValue(OPT_init, "_init");
   config->ltoAAPipeline = args.getLastArgValue(OPT_lto_aa_pipeline);
   config->ltoCSProfileGenerate = args.hasArg(OPT_lto_cs_profile_generate);
   config->ltoCSProfileFile = args.getLastArgValue(OPT_lto_cs_profile_file);
   config->ltoPGOWarnMismatch = args.hasFlag(OPT_lto_pgo_warn_mismatch,
                                             OPT_no_lto_pgo_warn_mismatch, true);
   config->ltoDebugPassManager = args.hasArg(OPT_lto_debug_pass_manager);
   config->ltoEmitAsm = args.hasArg(OPT_lto_emit_asm);
+  config->ltoExitOn = args.getLastArgValue(OPT_lto_exit_on_eq);
+  if (!config->ltoExitOn.empty() &&
+      !llvm::is_contained(ltoExitOnValues, config->ltoExitOn)) {
+    error("unknown --lto-exit-on value: " + config->ltoExitOn);
+    config->ltoExitOn = "";
+  }
   config->ltoNewPmPasses = args.getLastArgValue(OPT_lto_newpm_passes);
   config->ltoWholeProgramVisibility =
       args.hasFlag(OPT_lto_whole_program_visibility,
                    OPT_no_lto_whole_program_visibility, false);
   config->ltoo = args::getInteger(args, OPT_lto_O, 2);
   config->ltoObjPath = args.getLastArgValue(OPT_lto_obj_path_eq);
   config->ltoPartitions = args::getInteger(args, OPT_lto_partitions, 1);
   config->ltoSampleProfile = args.getLastArgValue(OPT_lto_sample_profile);
   config->ltoBasicBlockSections =
void LinkerDriver::link(opt::InputArgList &args) {
 
   // Skip the normal linked output if some LTO options are specified.
   //
   // For --thinlto-index-only, index file creation is performed in
   // compileBitcodeFiles, so we are done afterwards. --plugin-opt=emit-llvm and
   // --plugin-opt=emit-asm create output files in bitcode or assembly code,
   // respectively. When only certain thinLTO modules are specified for
   // compilation, the intermediate object file are the expected output.
-  const bool skipLinkedOutput = config->thinLTOIndexOnly || config->emitLLVM ||
-                                config->ltoEmitAsm ||
-                                !config->thinLTOModulesToCompile.empty();
+  // --lto-exit-on is specified to prevent the linker from executing after
+  // certain LTO stages.
+  const bool skipLinkedOutput =
+      config->thinLTOIndexOnly || config->emitLLVM || config->ltoEmitAsm ||
+      !config->ltoExitOn.empty() || !config->thinLTOModulesToCompile.empty();
 
   // Do link-time optimization if given files are LLVM bitcode files.
   // This compiles bitcode files into real object files.
   //
   // With this the symbol table should be complete. After this, no new names
   // except a few linker-synthesized ones will be added to the symbol table.
   const size_t numObjsBeforeLTO = ctx->objectFiles.size();
   invokeELFT(compileBitcodeFiles, skipLinkedOutput);

lld/ELF/LTO.cpp

static lto::Config createConfig() {
 
   if (config->ltoEmitAsm)
     c.CGFileType = CGFT_AssemblyFile;
 
   if (!config->saveTempsArgs.empty())
     checkError(c.addSaveTemps(config->outputFile.str() + ".",
                               /*UseInputModulePath*/ true,
                               config->saveTempsArgs));
+
+  if (!config->ltoExitOn.empty() && config->ltoExitOn != "prelink")
+    checkError(c.addExitOn(config->ltoExitOn));
   return c;
 }
 
 BitcodeCompiler::BitcodeCompiler() {
   // Initialize indexFile.
   if (!config->thinLTOIndexOnlyArg.empty())
     indexFile = openFile(config->thinLTOIndexOnlyArg);
 

lld/ELF/Options.td

 
 // LTO-related options.
 def lto_aa_pipeline: JJ<"lto-aa-pipeline=">,
   HelpText<"AA pipeline to run during LTO. Used in conjunction with -lto-newpm-passes">;
 def lto_debug_pass_manager: FF<"lto-debug-pass-manager">,
   HelpText<"Debug new pass manager">;
 def lto_emit_asm: FF<"lto-emit-asm">,
   HelpText<"Emit assembly code">;
+def lto_exit_on_eq: JJ<"lto-exit-on=">,
+  HelpText<"Stop the linker at a given stage. If using ThinLTO, linker will only exit after the given ThinLTO stage.">,
+  Values<"preopt,promote,internalize,import,opt,precodegen,prelink">;
 def no_lto_legacy_pass_manager: FF<"no-lto-legacy-pass-manager">,
   HelpText<"Use the new pass manager in LLVM">;
 def lto_newpm_passes: JJ<"lto-newpm-passes=">,
   HelpText<"Passes to run during LTO">;
 def lto_O: JJ<"lto-O">, MetaVarName<"<opt-level>">,
   HelpText<"Optimization level for LTO">;
 def lto_partitions: JJ<"lto-partitions=">,
   HelpText<"Number of LTO codegen partitions">;

lld/test/ELF/lto/lto-exit-on.ll

+;; This test is similar to llvm/test/ThinLTO/X86/exit-on.ll
+
+; REQUIRES: x86
+; UNSUPPORTED: system-windows
+;; Unsupported on Windows due to difficulty with escaping "opt" across platforms.
+;; lit substitutes 'opt' with /path/to/opt.
+
+; RUN: rm -fr %t && mkdir %t && cd %t
+; RUN: mkdir all all2 build
+; RUN: cd build
+
+; RUN: opt -thinlto-bc -o main.o %s
+; RUN: opt -thinlto-bc -o thin1.o %S/Inputs/thinlto.ll
+
+;; Check preopt
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=preopt
+; RUN: ls *.0.preopt*
+; RUN: not ls *.1.promote*
+; RUN: not ls *.2.internalize*
+; RUN: not ls *.3.import*
+; RUN: not ls *.4.opt*
+; RUN: not ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check promote
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=promote
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+;; 1 file is expected due to full LTO outputting .internalize bc
+;; before the .promote hook (unique to ThinLTO) is hit
+; RUN: ls *.2.internalize* | count 1
+; RUN: not ls *.3.import*
+; RUN: not ls *.4.opt*
+; RUN: not ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check internalize
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=internalize
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+;; 3 files are expected here and beyond due to full LTO outputting .internalize bc
+;; and the ThinLTO internalize hook outputting 2 more files
+; RUN: ls *.2.internalize* | count 3
+; RUN: not ls *.3.import*
+; RUN: not ls *.4.opt*
+; RUN: not ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check import
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=import
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+; RUN: ls *.2.internalize* | count 3
+; RUN: ls *.3.import*
+; RUN: not ls *.4.opt*
+; RUN: not ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check opt
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=\opt
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+; RUN: ls *.2.internalize* | count 3
+; RUN: ls *.3.import*
+; RUN: ls *.4.opt*
+; RUN: not ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check precodegen
+; RUN: rm -f *.o.*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=precodegen
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+; RUN: ls *.2.internalize* | count 3
+; RUN: ls *.3.import*
+; RUN: ls *.4.opt*
+; RUN: ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check prelink
+; RUN: rm -f *.o.* a.out*
+; RUN: ld.lld main.o thin1.o --save-temps --lto-exit-on=prelink
+; RUN: ls *.0.preopt*
+; RUN: ls *.1.promote*
+; RUN: ls *.2.internalize* | count 3
+; RUN: ls *.3.import*
+; RUN: ls *.4.opt*
+; RUN: ls *.5.precodegen*
+; RUN: ls a.out*.lto.o
+; RUN: not ls a.out
+
+;; Check output files are as expected
+; RUN: mv *.o.* a.out* %t/all2
+
+;; Create the .all dir with save-temps, without exit-on, then diff
+; RUN: ld.lld main.o thin1.o --save-temps
+; RUN: rm a.out
+; RUN: mv *.o.* a.out* %t/all
+; RUN: diff -r %t/all %t/all2
+
+;; Check input validation
+; RUN: not ld.lld main.o thin1.o --save-temps --lto-exit-on=notastage 2>&1 \
+; RUN: | FileCheck %s --check-prefix=ERR1
+; ERR1: unknown --lto-exit-on value: notastage
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @g()
+
+define i32 @_start() {
+  call void @g()
+  ret i32 0
+}

llvm/include/llvm/LTO/Config.h

struct Config {
   /// is prefixed by the given output file name and sets ResolutionFile to its
   /// file handle.
   ///
   /// SaveTempsArgs can be specified to select which temps to save.
   /// If SaveTempsArgs is not provided, all temps are saved.
   Error addSaveTemps(std::string OutputFileName,
                      bool UseInputModulePath = false,
                      const DenseSet<StringRef> &SaveTempsArgs = {});
+
+  /// Configures this Config object to stop executing further LTO/linker stages
+  /// upon reaching a specified stage.
+  /// This is done by modifying the above module hooks such that they:
+  /// 1) Run the existing hook, if any.
+  /// 2) return false.
+  Error addExitOn(StringRef LTOExitOn);
 };
 
 struct LTOLLVMDiagnosticHandler : public DiagnosticHandler {
   DiagnosticHandlerFunction *Fn;
   LTOLLVMDiagnosticHandler(DiagnosticHandlerFunction *DiagHandlerFn)
       : Fn(DiagHandlerFn) {}
   bool handleDiagnostics(const DiagnosticInfo &DI) override {
     (*Fn)(DI);

llvm/lib/LTO/LTOBackend.cpp

 #include "llvm/IR/Verifier.h"
 #include "llvm/LTO/LTO.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/ModuleSymbolTable.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/raw_ostream.h"
if (SaveTempsArgs.contains("precodegen"))
       setHook("5.precodegen", PreCodeGenModuleHook);
     if (SaveTempsArgs.contains("combinedindex"))
       CombinedIndexHook = SaveCombinedIndex;
   }
 
   return Error::success();
 }
 
+Error Config::addExitOn(StringRef LTOExitOn) {
+  auto setHook = [&](ModuleHookFn &Hook) {
+    // Keep track of the hook provided by the linker, which needs to run first.
+    ModuleHookFn LinkerHook = Hook;
+    Hook = [=](unsigned Task, const Module &M) {
+      if (LinkerHook)
+        LinkerHook(Task, M);
+      return false;
+    };
+  };
+
+  if (LTOExitOn == "preopt")
+    setHook(PreOptModuleHook);
+  else if (LTOExitOn == "promote")
+    setHook(PostPromoteModuleHook);
+  else if (LTOExitOn == "internalize")
+    setHook(PostInternalizeModuleHook);
+  else if (LTOExitOn == "import")
+    setHook(PostImportModuleHook);
+  else if (LTOExitOn == "opt")
+    setHook(PostOptModuleHook);
+  else if (LTOExitOn == "precodegen")
+    setHook(PreCodeGenModuleHook);
+  else
+    return make_error<StringError>("invalid addExitOn parameter: " + LTOExitOn,
+                                   errc::invalid_argument);
+
+  return Error::success();
+}
+
 #define HANDLE_EXTENSION(Ext)                                                  \
   llvm::PassPluginLibraryInfo get##Ext##PluginInfo();
 #include "llvm/Support/Extension.def"
 
 static void RegisterPassPlugins(ArrayRef<std::string> PassPlugins,
                                 PassBuilder &PB) {
 #define HANDLE_EXTENSION(Ext)                                                  \
   get##Ext##PluginInfo().RegisterPassBuilderCallbacks(PB);

llvm/test/ThinLTO/X86/exit-on.ll

+; UNSUPPORTED: system-windows
+;; Unsupported on Windows due to difficulty with escaping "opt" across platforms.
+;; lit substitutes 'opt' with /path/to/opt.
+
+; RUN: rm -rf %t && mkdir %t && cd %t
+
+;; Copy IR from import-constant.ll since it generates all the temps
+; RUN: opt -thinlto-bc %s -o 1.bc
+; RUN: opt -thinlto-bc %p/Inputs/import-constant.ll -o 2.bc
+; RUN: mkdir all build
+
+;; Check preopt
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=preopt
+; RUN: ls build/*.0.preopt*
+; RUN: not ls build/*.1.promote*
+; RUN: not ls build/*.2.internalize*
+; RUN: not ls build/*.3.import*
+; RUN: not ls build/*.4.opt*
+; RUN: not ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Check promote
+; RUN: rm -f build/*
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=promote
+; RUN: ls build/*.0.preopt*
+; RUN: ls build/*.1.promote*
+;; 1 file is expected due to full LTO outputting .internalize bc
+;; before the .promote hook (unique to ThinLTO) is hit
+; RUN: ls build/*.2.internalize* | count 1
+; RUN: not ls build/*.3.import*
+; RUN: not ls build/*.4.opt*
+; RUN: not ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Check internalize
+; RUN: rm -f build/*
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=internalize
+; RUN: ls build/*.0.preopt*
+; RUN: ls build/*.1.promote*
+;; 3 files are expected here and beyond due to full LTO outputting .internalize bc
+;; and the ThinLTO internalize hook outputting 2 more files
+; RUN: ls build/*.2.internalize* | count 3
+; RUN: not ls build/*.3.import*
+; RUN: not ls build/*.4.opt*
+; RUN: not ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Check import
+; RUN: rm -f build/*
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=import
+; RUN: ls build/*.0.preopt*
+; RUN: ls build/*.1.promote*
+; RUN: ls build/*.2.internalize* | count 3
+; RUN: ls build/*.3.import*
+; RUN: not ls build/*.4.opt*
+; RUN: not ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Check opt
+; RUN: rm -f build/*
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=\opt
+; RUN: ls build/*.0.preopt*
+; RUN: ls build/*.1.promote*
+; RUN: ls build/*.2.internalize* | count 3
+; RUN: ls build/*.3.import*
+; RUN: ls build/*.4.opt*
+; RUN: not ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Check precodegen
+; RUN: llvm-lto2 run 1.bc 2.bc -o build/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps -exit-on=precodegen
+; RUN: ls build/*.0.preopt*
+; RUN: ls build/*.1.promote*
+; RUN: ls build/*.2.internalize* | count 3
+; RUN: ls build/*.3.import*
+; RUN: ls build/*.4.opt*
+; RUN: ls build/*.5.precodegen*
+; RUN: not ls build/a.out.1
+; RUN: not ls build/a.out.2
+
+;; Create the .all dir with save-temps, without exit-on, then diff
+; RUN: llvm-lto2 run 1.bc 2.bc -o all/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -save-temps
+; RUN: rm all/a.out.1 all/a.out.2
+; RUN: diff -r all build
+
+;; Check error message
+; RUN: not llvm-lto2 run 1.bc 2.bc -o all/a.out \
+; RUN:    -import-constants-with-refs -r=1.bc,main,plx -r=1.bc,_Z6getObjv,l \
+; RUN:    -r=2.bc,_Z6getObjv,pl -r=2.bc,val,pl -r=2.bc,outer,pl \
+; RUN:    -exit-on=prelink 2>&1 \
+; RUN: | FileCheck %s --check-prefix=ERR1
+; ERR1: invalid addExitOn parameter: prelink
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.S = type { i32, i32, i32* }
+
+define dso_local i32 @main() local_unnamed_addr {
+entry:
+  %call = tail call %struct.S* @_Z6getObjv()
+  %d = getelementptr inbounds %struct.S, %struct.S* %call, i64 0, i32 0
+  %0 = load i32, i32* %d, align 8
+  %v = getelementptr inbounds %struct.S, %struct.S* %call, i64 0, i32 1
+  %1 = load i32, i32* %v, align 4
+  %add = add nsw i32 %1, %0
+  ret i32 %add
+}
+
+declare dso_local %struct.S* @_Z6getObjv() local_unnamed_addr

llvm/tools/llvm-lto2/llvm-lto2.cpp

 static cl::opt<std::string> OptPipeline("opt-pipeline",
                                         cl::desc("Optimizer Pipeline"),
                                         cl::value_desc("pipeline"));
 
 static cl::opt<std::string> AAPipeline("aa-pipeline",
                                        cl::desc("Alias Analysis Pipeline"),
                                        cl::value_desc("aapipeline"));
 
+static cl::opt<std::string> ExitOn(
+    "exit-on", cl::desc("Stop LTO at a given stage"),
+    cl::value_desc("One of: preopt,promote,internalize,import,opt,precodegen"));
+
 static cl::opt<bool> SaveTemps("save-temps", cl::desc("Save temporary files"));
 
 static cl::list<std::string> SelectSaveTemps(
     "select-save-temps",
     cl::value_desc("One, or multiple of: "
                    "resolution,preopt,promote,internalize,import,opt,precodegen"
                    ",combinedindex"),
     cl::desc("Save selected temporary files. Cannot be specified together with "
for (auto &S : SelectSaveTemps)
       else {
         llvm::errs() << ("invalid -select-save-temps argument: " + S) << '\n';
         return 1;
       }
     check(Conf.addSaveTemps(OutputFilename + ".", false, SaveTempsArgs),
           "Config::addSaveTemps failed");
   }
 
+  if (!ExitOn.empty())
+    check(Conf.addExitOn(ExitOn), "Config::addExitOn failed");
+
   // Optimization remarks.
   Conf.RemarksFilename = RemarksFilename;
   Conf.RemarksPasses = RemarksPasses;
   Conf.RemarksWithHotness = RemarksWithHotness;
   Conf.RemarksHotnessThreshold = RemarksHotnessThreshold;
   Conf.RemarksFormat = RemarksFormat;
 
   Conf.SampleProfile = SamplePGOFile;