From 42fead6834ef3a96c765ea545c2d2bac951f7e98 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin@nvidia.com>
Date: Mon, 3 Oct 2022 16:54:12 -0700
Subject: [PATCH] [flang][tco] Engineering option for running only CodeGen
 passes.

This option allows running only CodeGen passes and then translating
FIR to LLVM IR. I am using it to fetch optimized FIR after bbc,
hand-modify it and then produce LLVM IR that can be fed to clang.
---
 flang/test/Driver/tco-code-gen-llvm.fir | 13 +++++++++++++
 flang/tools/tco/tco.cpp                 | 16 +++++++++++++---
 2 files changed, 26 insertions(+), 3 deletions(-)
 create mode 100644 flang/test/Driver/tco-code-gen-llvm.fir

diff --git a/flang/test/Driver/tco-code-gen-llvm.fir b/flang/test/Driver/tco-code-gen-llvm.fir
new file mode 100644
index 000000000000..f0641865c09b
--- /dev/null
+++ b/flang/test/Driver/tco-code-gen-llvm.fir
@@ -0,0 +1,13 @@
+// Test the tco's -code-gen-llvm option.
+
+// RUN: tco -code-gen-llvm %s 2>&1 | FileCheck %s
+
+// Check that FIR is translated into LLVM IR, and that
+// there is no any FIR output.
+
+// CHECK-NOT: func.func
+// CHECK: define void @_QPfoo
+// CHECK-NOT: func.func
+func.func @_QPfoo() {
+  return
+}
diff --git a/flang/tools/tco/tco.cpp b/flang/tools/tco/tco.cpp
index 5fe0d8114742..a1b60aa7443a 100644
--- a/flang/tools/tco/tco.cpp
+++ b/flang/tools/tco/tco.cpp
@@ -53,6 +53,11 @@ static cl::opt<std::string> targetTriple("target",
                                          cl::desc("specify a target triple"),
                                          cl::init("native"));
 
+static cl::opt<bool> codeGenLLVM(
+    "code-gen-llvm",
+    cl::desc("Run only CodeGen passes and translate FIR to LLVM IR"),
+    cl::init(false));
+
 #include "flang/Tools/CLOptions.inc"
 
 static void printModuleBody(mlir::ModuleOp mod, raw_ostream &output) {
@@ -112,15 +117,20 @@ compileFIR(const mlir::PassPipelineCLParser &passPipeline) {
     if (mlir::failed(passPipeline.addToPipeline(pm, errorHandler)))
       return mlir::failure();
   } else {
-    // Run tco with O2 by default.
-    fir::createMLIRToLLVMPassPipeline(pm, llvm::OptimizationLevel::O2);
+    if (codeGenLLVM) {
+      // Run only CodeGen passes.
+      fir::createDefaultFIRCodeGenPassPipeline(pm);
+    } else {
+      // Run tco with O2 by default.
+      fir::createMLIRToLLVMPassPipeline(pm, llvm::OptimizationLevel::O2);
+    }
     fir::addLLVMDialectToLLVMPass(pm, out.os());
   }
 
   // run the pass manager
   if (mlir::succeeded(pm.run(*owningRef))) {
     // passes ran successfully, so keep the output
-    if (emitFir || passPipeline.hasAnyOccurrences())
+    if ((emitFir || passPipeline.hasAnyOccurrences()) && !codeGenLLVM)
       printModuleBody(*owningRef, out.os());
     out.keep();
     return mlir::success();