[llvm] [LTO] Add an option to select a pre-link pipeline (PR #82585)

[Igor Kudrin via llvm-commits]

https://github.com/igorkudrin created https://github.com/llvm/llvm-project/pull/82585

The option specifies the variant of the optimization pipeline to be used in the LTO backend. The pre-link pipeline can be useful in cases where LTO is used in an intermediate step that combines and optimizes a subset of input files, and the results are used later in the final link.

This implements an idea suggested in https://github.com/llvm/llvm-project/pull/71268#issuecomment-1955086700.

>From e01980ae6037c62f4b98dfbcd12900581f08a23a Mon Sep 17 00:00:00 2001
From: Igor Kudrin <ikudrin at accesssoftek.com>
Date: Wed, 21 Feb 2024 21:47:25 -0800
Subject: [PATCH] [LTO] Add an option to select a pre-link pipeline

The option specifies the variant of the optimization pipeline to be used
in the LTO backend. The pre-link pipeline can be useful in cases where
LTO is used in an intermediate step that combines and optimizes a subset
of input files, and the results are used later in the final link.
---
 llvm/include/llvm/LTO/Config.h            |  3 +++
 llvm/lib/LTO/LTOBackend.cpp               |  8 ++++++--
 llvm/test/tools/llvm-lto2/X86/pipeline.ll | 13 +++++++++++++
 llvm/tools/llvm-lto2/llvm-lto2.cpp        |  4 ++++
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/LTO/Config.h b/llvm/include/llvm/LTO/Config.h
index 6fb55f1cf1686a..9a87e8cb309ef3 100644
--- a/llvm/include/llvm/LTO/Config.h
+++ b/llvm/include/llvm/LTO/Config.h
@@ -63,6 +63,9 @@ struct Config {
   /// Use the standard optimization pipeline.
   bool UseDefaultPipeline = false;
 
+  /// Use pre-link pipeline.
+  bool UsePreLinkPipeline = false;
+
   /// Flag to indicate that the optimizer should not assume builtins are present
   /// on the target.
   bool Freestanding = false;
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
index 7b3a7590dfa743..31ca49d02065c1 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -333,9 +333,13 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   } else if (Conf.UseDefaultPipeline) {
     MPM.addPass(PB.buildPerModuleDefaultPipeline(OL));
   } else if (IsThinLTO) {
-    MPM.addPass(PB.buildThinLTODefaultPipeline(OL, ImportSummary));
+    MPM.addPass(Conf.UsePreLinkPipeline
+                    ? PB.buildThinLTOPreLinkDefaultPipeline(OL)
+                    : PB.buildThinLTODefaultPipeline(OL, ImportSummary));
   } else {
-    MPM.addPass(PB.buildLTODefaultPipeline(OL, ExportSummary));
+    MPM.addPass(Conf.UsePreLinkPipeline
+                    ? PB.buildLTOPreLinkDefaultPipeline(OL)
+                    : PB.buildLTODefaultPipeline(OL, ExportSummary));
   }
 
   if (!Conf.DisableVerify)
diff --git a/llvm/test/tools/llvm-lto2/X86/pipeline.ll b/llvm/test/tools/llvm-lto2/X86/pipeline.ll
index 93f30d0ee61bce..774405f08385fa 100644
--- a/llvm/test/tools/llvm-lto2/X86/pipeline.ll
+++ b/llvm/test/tools/llvm-lto2/X86/pipeline.ll
@@ -40,3 +40,16 @@ define void @patatino() {
 ; RUN:  FileCheck %s --check-prefix=AAERR
 
 ; AAERR: LLVM ERROR: unable to parse AA pipeline description 'patatino': unknown alias analysis name 'patatino'
+
+;; Check that LTO can be configured to use the pre-link pipeline.
+; RUN: opt -module-summary %s -o %tthin.bc
+; RUN: llvm-lto2 run %t1.bc -o %t.o -r %t1.bc,patatino,px \
+; RUN:   -prelink-pipeline -debug-pass-manager 2>&1 | \
+; RUN:   FileCheck %s --check-prefix=PRELINK
+; RUN: llvm-lto2 run %tthin.bc -o %t.o -r %tthin.bc,patatino,px \
+; RUN:   -prelink-pipeline -debug-pass-manager 2>&1 | \
+; RUN:   FileCheck %s --check-prefix=PRELINK
+
+; PRELINK-NOT: EliminateAvailableExternallyPass
+; PRELINK:     Running pass: InlinerPass on (patatino)
+; PRELINK-NOT: EliminateAvailableExternallyPass
diff --git a/llvm/tools/llvm-lto2/llvm-lto2.cpp b/llvm/tools/llvm-lto2/llvm-lto2.cpp
index d5de4f6b1a277c..7c574d361400ee 100644
--- a/llvm/tools/llvm-lto2/llvm-lto2.cpp
+++ b/llvm/tools/llvm-lto2/llvm-lto2.cpp
@@ -64,6 +64,9 @@ static cl::opt<std::string> AAPipeline("aa-pipeline",
                                        cl::desc("Alias Analysis Pipeline"),
                                        cl::value_desc("aapipeline"));
 
+static cl::opt<bool> PreLinkPipeline("prelink-pipeline",
+                                     cl::desc("Use the pre-link pipeline"));
+
 static cl::opt<bool> SaveTemps("save-temps", cl::desc("Save temporary files"));
 
 static cl::list<std::string> SelectSaveTemps(
@@ -315,6 +318,7 @@ static int run(int argc, char **argv) {
   // Run a custom pipeline, if asked for.
   Conf.OptPipeline = OptPipeline;
   Conf.AAPipeline = AAPipeline;
+  Conf.UsePreLinkPipeline = PreLinkPipeline;
 
   Conf.OptLevel = OptLevel - '0';
   Conf.Freestanding = EnableFreestanding;