From 75a5eaf7c6d60e95b11ea572b78fdb8788d15ddc Mon Sep 17 00:00:00 2001
From: Shilei Tian <i@tianshilei.me>
Date: Mon, 20 Dec 2021 20:01:37 -0500
Subject: [PATCH] [LTO] Add a function `LTOCodeGenerator::getMergedModule`

One of the uses of `LTOCodeGenerator` is to take it as a middle+back end. Sometimes
it is very helpful to access, especially get information from the optimized module.
If the information can be changed in optimization, it cannot be get before the
module is added to `LTOCodeGenerator`. This patch adds a function
`LTOCodeGenerator::getMergedModule` to access the `MergedModule`.

Reviewed By: steven_wu

Differential Revision: https://reviews.llvm.org/D114201
---
 llvm/include/llvm/LTO/legacy/LTOCodeGenerator.h |  2 ++
 llvm/test/tools/llvm-lto/dump.ll                | 14 ++++++++++++++
 llvm/tools/llvm-lto/llvm-lto.cpp                | 17 +++++++++++++++--
 3 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/tools/llvm-lto/dump.ll

diff --git a/llvm/include/llvm/LTO/legacy/LTOCodeGenerator.h b/llvm/include/llvm/LTO/legacy/LTOCodeGenerator.h
index 333f483f29c5..91c1decac373 100644
--- a/llvm/include/llvm/LTO/legacy/LTOCodeGenerator.h
+++ b/llvm/include/llvm/LTO/legacy/LTOCodeGenerator.h
@@ -193,6 +193,8 @@ struct LTOCodeGenerator {
   void resetMergedModule() { MergedModule.reset(); }
   void DiagnosticHandler(const DiagnosticInfo &DI);
 
+  Module &getMergedModule() { return *MergedModule; }
+
 private:
   /// Verify the merged module on first call.
   ///
diff --git a/llvm/test/tools/llvm-lto/dump.ll b/llvm/test/tools/llvm-lto/dump.ll
new file mode 100644
index 000000000000..57e6b757664f
--- /dev/null
+++ b/llvm/test/tools/llvm-lto/dump.ll
@@ -0,0 +1,14 @@
+; RUN: llvm-as < %s > %t1
+; RUN: llvm-lto --dump-linked-module %t1 2>&1 | FileCheck %s
+
+; CHEKCK: target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; CHECK: define void @f() {
+define void @f() {
+; CHECK-NEXT: entry:
+entry:
+; CHECK-NEXT: ret void
+  ret void
+; CHECK-NEXT: }
+}
diff --git a/llvm/tools/llvm-lto/llvm-lto.cpp b/llvm/tools/llvm-lto/llvm-lto.cpp
index 995ebacacb87..be91ec25d7ad 100644
--- a/llvm/tools/llvm-lto/llvm-lto.cpp
+++ b/llvm/tools/llvm-lto/llvm-lto.cpp
@@ -197,6 +197,11 @@ static cl::opt<bool>
                    cl::desc("Write merged LTO module to file before CodeGen"),
                    cl::cat(LTOCategory));
 
+static cl::opt<bool>
+    DumpLinkedModule("dump-linked-module", cl::init(false),
+                     cl::desc("Dump linked LTO module before optimize"),
+                     cl::cat(LTOCategory));
+
 static cl::list<std::string> InputFilenames(cl::Positional, cl::OneOrMore,
                                             cl::desc("<input bitcode files>"),
                                             cl::cat(LTOCategory));
@@ -316,11 +321,11 @@ namespace {
       if (!CurrentActivity.empty())
         OS << ' ' << CurrentActivity;
       OS << ": ";
-  
+
       DiagnosticPrinterRawOStream DP(OS);
       DI.print(DP);
       OS << '\n';
-  
+
       if (DI.getSeverity() == DS_Error)
         exit(1);
       return true;
@@ -1115,6 +1120,14 @@ int main(int argc, char **argv) {
       error("error compiling the code");
 
   } else {
+    if (DumpLinkedModule) {
+      if (SaveLinkedModuleFile)
+        error(": -dump-linked-module must be used without -dump-linked-module");
+
+      Module &M = CodeGen.getMergedModule();
+      M.dump();
+    }
+
     if (Parallelism != 1)
       error("-j must be specified together with -o");