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

[Igor Kudrin via llvm-commits]

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

>From be31c3a70b1910e203a8d45b2ee082124f670d19 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 1/2] [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 482b6e55a19d35..6049ca74a54fca 100644
--- a/llvm/include/llvm/LTO/Config.h
+++ b/llvm/include/llvm/LTO/Config.h
@@ -60,6 +60,9 @@ struct Config {
   bool VerifyEach = false;
   bool DisableVerify = 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 6cfe67779b1a7d..5fbd0f90b97c7a 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -331,9 +331,13 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
                          Conf.OptPipeline + "': " + toString(std::move(Err)));
     }
   } 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;

>From 98bb6618a80599c2f239a04a1297e24bcc4e794e Mon Sep 17 00:00:00 2001
From: Igor Kudrin <ikudrin at accesssoftek.com>
Date: Thu, 22 Feb 2024 11:04:04 -0800
Subject: [PATCH 2/2] fixup

---
 llvm/include/llvm/LTO/Config.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/include/llvm/LTO/Config.h b/llvm/include/llvm/LTO/Config.h
index 6049ca74a54fca..87f2cc5108d80e 100644
--- a/llvm/include/llvm/LTO/Config.h
+++ b/llvm/include/llvm/LTO/Config.h
@@ -60,7 +60,7 @@ struct Config {
   bool VerifyEach = false;
   bool DisableVerify = false;
 
-  /// Use pre-link pipeline.
+  /// Use pre-link optimization pipeline.
   bool UsePreLinkPipeline = false;
 
   /// Flag to indicate that the optimizer should not assume builtins are present