[flang] Add -O flag to tco

[parabola94]

At the moment, there is no established way to emit LLVM IR before optimization in LLVM. tco can partially does, but the optimization level is fixed to 2. This patch adds -O flag to tco.

Note that this is not completely equivalent to the frontend option. tco does not accept -O flag without numbers, and the default optimization level is not 0 but 2.

[llvmbot]

@llvm/pr-subscribers-flang-driver

Author: None (parabola94)

Changes

At the moment, there is no established way to emit LLVM IR before optimization in LLVM. tco can partially does, but the optimization level is fixed to 2. This patch adds -O flag to tco.

Note that this is not completely equivalent to the frontend option. tco does not accept -O flag without numbers, and the default optimization level is not 0 but 2.

---

Full diff: https://github.com/llvm/llvm-project/pull/151869.diff

2 Files Affected:

• (modified) flang/test/Driver/tco-test-gen.fir (+5-4)
• (modified) flang/tools/tco/tco.cpp (+33-2)

diff --git a/flang/test/Driver/tco-test-gen.fir b/flang/test/Driver/tco-test-gen.fir
index 38d4e50ecf3aa..ec566488e32dd 100644
--- a/flang/test/Driver/tco-test-gen.fir
+++ b/flang/test/Driver/tco-test-gen.fir
@@ -1,5 +1,6 @@
 // RUN: tco -emit-final-mlir %s | FileCheck %s --check-prefixes=CHECK,AA,CMPLX
 // RUN: tco -emit-final-mlir -enable-aa=false %s | FileCheck %s --check-prefixes=CHECK,NOAA,CMPLX
+// RUN: tco -emit-final-mlir -O0 -enable-aa %s | FileCheck %s --check-prefixes=CHECK,NOAA,CMPLX
 // RUN: tco -emit-final-mlir -simplify-mlir %s | FileCheck %s --check-prefixes=CHECK,AA,SIMPLE
 // RUN: tco -emit-final-mlir -enable-aa=false -simplify-mlir %s | FileCheck %s --check-prefixes=CHECK,NOAA,SIMPLE
 // RUN: tco -test-gen %s | FileCheck %s --check-prefixes=CHECK,NOAA,SIMPLE
@@ -37,10 +38,10 @@ func.func @_QPtest(%arg0: !fir.ref<i32> {fir.bindc_name = "num"}, %arg1: !fir.re
 }
 
 // CHECK-LABEL:   llvm.func @_QPtest(
-// CHECK-SAME:      %[[ARG0:.*]]: !llvm.ptr {fir.bindc_name = "num", llvm.nocapture},
-// CHECK-SAME:      %[[ARG1:.*]]: !llvm.ptr {fir.bindc_name = "lb", llvm.nocapture},
-// CHECK-SAME:      %[[ARG2:.*]]: !llvm.ptr {fir.bindc_name = "ub", llvm.nocapture},
-// CHECK-SAME:      %[[ARG3:.*]]: !llvm.ptr {fir.bindc_name = "step", llvm.nocapture}) {
+// CHECK-SAME:      %[[ARG0:.*]]: !llvm.ptr {fir.bindc_name = "num"{{[^}]*}}},
+// CHECK-SAME:      %[[ARG1:.*]]: !llvm.ptr {fir.bindc_name = "lb"{{[^}]*}}},
+// CHECK-SAME:      %[[ARG2:.*]]: !llvm.ptr {fir.bindc_name = "ub"{{[^}]*}}},
+// CHECK-SAME:      %[[ARG3:.*]]: !llvm.ptr {fir.bindc_name = "step"{{[^}]*}}}) {
 
 // CMPLX:           %[[VAL_0:.*]] = llvm.mlir.constant(1 : i64) : i64
 // CMPLX:           %[[VAL_1:.*]] = llvm.alloca %[[VAL_0]] x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
diff --git a/flang/tools/tco/tco.cpp b/flang/tools/tco/tco.cpp
index d8daf8769cb21..b6d7c142c19df 100644
--- a/flang/tools/tco/tco.cpp
+++ b/flang/tools/tco/tco.cpp
@@ -51,6 +51,12 @@ static cl::opt<bool> emitFir("emit-fir",
                              cl::desc("Parse and pretty-print the input"),
                              cl::init(false));
 
+static cl::opt<unsigned>
+    OptLevel("O",
+             cl::desc("Optimization level. [-O0, -O1, -O2, or -O3] "
+                      "(default = '-O2')"),
+             cl::Prefix, cl::init(2));
+
 static cl::opt<std::string> targetTriple("target",
                                          cl::desc("specify a target triple"),
                                          cl::init("native"));
@@ -96,6 +102,22 @@ static void printModule(mlir::ModuleOp mod, raw_ostream &output) {
   output << mod << '\n';
 }
 
+static std::optional<llvm::OptimizationLevel>
+getOptimizationLevel(unsigned level) {
+  switch (level) {
+  default:
+    return std::nullopt;
+  case 0:
+    return llvm::OptimizationLevel::O0;
+  case 1:
+    return llvm::OptimizationLevel::O1;
+  case 2:
+    return llvm::OptimizationLevel::O2;
+  case 3:
+    return llvm::OptimizationLevel::O3;
+  }
+}
+
 // compile a .fir file
 static llvm::LogicalResult
 compileFIR(const mlir::PassPipelineCLParser &passPipeline) {
@@ -157,9 +179,18 @@ compileFIR(const mlir::PassPipelineCLParser &passPipeline) {
     if (mlir::failed(passPipeline.addToPipeline(pm, errorHandler)))
       return mlir::failure();
   } else {
-    MLIRToLLVMPassPipelineConfig config(llvm::OptimizationLevel::O2);
+    std::optional<llvm::OptimizationLevel> level =
+        getOptimizationLevel(OptLevel);
+    if (!level) {
+      errs() << "Error invalid optimization level\n";
+      return mlir::failure();
+    }
+    MLIRToLLVMPassPipelineConfig config(*level);
+    // TODO: config.StackArrays should be set here?
     config.EnableOpenMP = true;  // assume the input contains OpenMP
-    config.AliasAnalysis = enableAliasAnalysis && !testGeneratorMode;
+    config.AliasAnalysis =
+        OptLevel > 0 && enableAliasAnalysis && !testGeneratorMode;
+    config.LoopVersioning = OptLevel > 2;
     if (codeGenLLVM) {
       // Run only CodeGen passes.
       fir::createDefaultFIRCodeGenPassPipeline(pm, config);

[tblah]

Looks good to me. Thank you for contributing this.

[tblah]

ultra-nit: for me this is a convincing example that -enable-aa should take precidence over -O0 (and enable alias analysis even at -O0). But I don't think it matters too much for a developer tool.

[parabola94]

Thank you for your feedback. I agree with you. I've removed this.

[tblah]

stack arrays isn't enabled at -O3 (only -Ofast). I think it is okay to leave this as TODO for another patch if anyone wants a flag for it.

[parabola94]

@tblah Thank you for the review! Could you merge this on my behalf? I do not have write permissions.