ExecutionEngine OptUtils: support -On flags in string-based initialization

Original implementation of OutUtils provided two different LLVM IR module
transformers to be used with the MLIR ExecutionEngine: OptimizingTransformer
parameterized by the optimization levels (similar to -O3 flags) and
LLVMPassesTransformer parameterized by the string formatted similarly to
command line options of LLVM's "opt" tool without support for -O* flags.
Introduce such support by declaring the flags inside the parser and by
populating the pass managers similarly to what "opt" does.  Remove the
additional flags from mlir-cpu-runner as they can now be wrapped into
`-llvm-opts` together with other LLVM-related flags.

PiperOrigin-RevId: 236107292

mlir/lib/ExecutionEngine/OptUtils.cpp

@@ -33,15 +33,37 @@
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include <climits>
 #include <mutex>
 
+// A category for options that should be passed to -llvm-opts.
+static llvm::cl::OptionCategory
+    optFlags("opt-like flags (pass to -llvm-ops=\"\")");
+
+// LLVM pass configuration CLI flag.
+static llvm::cl::list<const llvm::PassInfo *, bool, llvm::PassNameParser>
+    llvmPasses(llvm::cl::desc("LLVM optimizing passes to run"),
+               llvm::cl::cat(optFlags));
+
+// CLI variables for -On options.
+static llvm::cl::opt<bool> optO0("O0", llvm::cl::desc("Run opt O0 passes"),
+                                 llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool> optO1("O1", llvm::cl::desc("Run opt O1 passes"),
+                                 llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool> optO2("O2", llvm::cl::desc("Run opt O2 passes"),
+                                 llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool> optO3("O3", llvm::cl::desc("Run opt O3 passes"),
+                                 llvm::cl::cat(optFlags));
+
 // Run the module and function passes managed by the module manager.
 static void runPasses(llvm::legacy::PassManager &modulePM,
                       llvm::legacy::FunctionPassManager &funcPM,
                       llvm::Module &m) {
+  funcPM.doInitialization();
   for (auto &func : m) {
     funcPM.run(func);
   }
+  funcPM.doFinalization();
   modulePM.run(m);
 }
 
@@ -61,35 +83,63 @@ void mlir::initializeLLVMPasses() {
   llvm::initializeVectorization(registry);
 }
 
+// Populate pass managers according to the optimization and size levels.
+// This behaves similarly to LLVM opt.
+static void populatePassManagers(llvm::legacy::PassManager &modulePM,
+                                 llvm::legacy::FunctionPassManager &funcPM,
+                                 unsigned optLevel, unsigned sizeLevel) {
+  llvm::PassManagerBuilder builder;
+  builder.OptLevel = optLevel;
+  builder.SizeLevel = sizeLevel;
+  builder.Inliner = llvm::createFunctionInliningPass(
+      optLevel, sizeLevel, /*DisableInlineHotCallSite=*/false);
+  builder.LoopVectorize = optLevel > 1 && sizeLevel < 2;
+  builder.SLPVectorize = optLevel > 1 && sizeLevel < 2;
+  builder.DisableUnrollLoops = (optLevel == 0);
+
+  builder.populateModulePassManager(modulePM);
+  builder.populateFunctionPassManager(funcPM);
+}
+
 // Create and return a lambda that uses LLVM pass manager builder to set up
 // optimizations based on the given level.
 std::function<llvm::Error(llvm::Module *)>
 mlir::makeOptimizingTransformer(unsigned optLevel, unsigned sizeLevel) {
   return [optLevel, sizeLevel](llvm::Module *m) -> llvm::Error {
-    llvm::PassManagerBuilder builder;
-    builder.OptLevel = optLevel;
-    builder.SizeLevel = sizeLevel;
-    builder.Inliner = llvm::createFunctionInliningPass(
-        optLevel, sizeLevel, /*DisableInlineHotCallSite=*/false);
 
     llvm::legacy::PassManager modulePM;
     llvm::legacy::FunctionPassManager funcPM(m);
-    builder.populateModulePassManager(modulePM);
-    builder.populateFunctionPassManager(funcPM);
+    populatePassManagers(modulePM, funcPM, optLevel, sizeLevel);
     runPasses(modulePM, funcPM, *m);
 
     return llvm::Error::success();
   };
 }
 
+// Check if the opt flag is set and if it was located before `pos`.  If so,
+// popuplate the module and the function pass managers with the passes
+// corresponding to the `level` of optimization and reset the flag.
+static void
+populatePassManagersOptLevel(llvm::cl::opt<bool> &opt, unsigned level,
+                             unsigned pos, llvm::legacy::PassManager &modulePM,
+                             llvm::legacy::FunctionPassManager &funcPM) {
+  if (opt && opt.getPosition() < pos) {
+    opt = false;
+    populatePassManagers(modulePM, funcPM, level, /*sizeLevel=*/0);
+  }
+}
+
 // Create and return a lambda that leverages LLVM PassInfo command line parser
 // to construct passes given the command line flags that come from the given
 // string rather than from the command line.
 std::function<llvm::Error(llvm::Module *)>
 mlir::makeLLVMPassesTransformer(std::string config) {
   return [config](llvm::Module *m) -> llvm::Error {
-    static llvm::cl::list<const llvm::PassInfo *, bool, llvm::PassNameParser>
-        llvmPasses(llvm::cl::desc("LLVM optimizing passes to run"));
+    // Storage for -On flags, the index in this array corresponds to the
+    // optimization level.  Do not add anything else.
+    llvm::SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
+        optO0, optO1, optO2, optO3};
+
     llvm::BumpPtrAllocator allocator;
     llvm::StringSaver saver(allocator);
     llvm::SmallVector<const char *, 16> args;
@@ -98,11 +148,19 @@ mlir::makeLLVMPassesTransformer(std::string config) {
     llvm::cl::ParseCommandLineOptions(args.size(), args.data());
 
     llvm::legacy::PassManager modulePM;
+    llvm::legacy::FunctionPassManager funcPM(m);
 
-    for (const auto *passInfo : llvmPasses) {
+    for (unsigned i = 0, e = llvmPasses.size(); i < e; ++i) {
+      const auto *passInfo = llvmPasses[i];
       if (!passInfo->getNormalCtor())
         continue;
 
+      // If there is any of -On flags textually before this pass flag, populate
+      // the pass managers with the corresponding passes and reset the flag.
+      for (unsigned j = 0; j < 4; ++j)
+        populatePassManagersOptLevel(
+            optFlags[j].get(), j, llvmPasses.getPosition(i), modulePM, funcPM);
+
       auto *pass = passInfo->createPass();
       if (!pass)
         return llvm::make_error<llvm::StringError>(
@@ -111,8 +169,20 @@ mlir::makeLLVMPassesTransformer(std::string config) {
 
       modulePM.add(pass);
     }
+    // Populate the pass managers with passes corresponding to the -On flags
+    // that have not been used yet.  Use UINT_MAX as the position index before
+    // which the -On flag should appear as an always-true marker.
+    for (unsigned j = 0; j < 4; ++j)
+      populatePassManagersOptLevel(optFlags[j].get(), j, /*pos=*/UINT_MAX,
+                                   modulePM, funcPM);
 
-    modulePM.run(*m);
+    // Run the -On function passes, then all the other passes.  Note that
+    // manually requested function passes were added to modulePM and will be
+    // executed in order with manual/-On module passes.  The function pass
+    // manager is only populated in reaction to -On flags with passes that are
+    // supposed to run "as soon as functions are created" according to the doc.
+    // This behavior is identical to that of LLVM's "opt" tool.
+    runPasses(modulePM, funcPM, *m);
     return llvm::Error::success();
   };
 }

mlir/test/mlir-cpu-runner/simple.mlir

@@ -1,7 +1,7 @@
 // RUN: mlir-cpu-runner %s | FileCheck %s
 // RUN: mlir-cpu-runner -e foo -init-value 1000 %s | FileCheck -check-prefix=NOMAIN %s
-// RUN: mlir-cpu-runner %s -O3 | FileCheck %s
-// RUN: mlir-cpu-runner %s -llvm-opts="-loop-distribute -loop-vectorize" | FileCheck %s
+// RUN: mlir-cpu-runner %s -llvm-opts="-O3" | FileCheck %s
+// RUN: mlir-cpu-runner %s -llvm-opts="-O3 -loop-distribute -loop-vectorize" | FileCheck %s
 
 func @fabsf(f32) -> f32
 

mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp

@@ -58,19 +58,6 @@ static llvm::cl::opt<std::string> llvmPasses(
     "llvm-opts",
     llvm::cl::desc("LLVM passes to run, syntax same as the opt tool"));
 
-static llvm::cl::opt<bool>
-    llvmO0("O0",
-           llvm::cl::desc("Optimization level 0, similar to LLVM opt -O0"));
-static llvm::cl::opt<bool>
-    llvmO1("O1",
-           llvm::cl::desc("Optimization level 1, similar to LLVM opt -O1"));
-static llvm::cl::opt<bool>
-    llvmO2("O2",
-           llvm::cl::desc("Optimization level 2, similar to LLVM opt -O2"));
-static llvm::cl::opt<bool>
-    llvmO3("O3",
-           llvm::cl::desc("Optimization level 3, similar to LLVM opt -O3"));
-
 static std::unique_ptr<Module> parseMLIRInput(StringRef inputFilename,
                                               MLIRContext *context) {
   // Set up the input file.
@@ -170,12 +157,6 @@ int main(int argc, char **argv) {
 
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
-  if ((llvmO0 || llvmO1 || llvmO2 || llvmO3) &&
-      !llvmPasses.getValue().empty()) {
-    llvm::errs() << "cannot use -O? together with -llvm-passes\n";
-    return EXIT_FAILURE;
-  }
-
   initializeLLVM();
   mlir::initializeLLVMPasses();
 
@@ -186,20 +167,7 @@ int main(int argc, char **argv) {
     return 1;
   }
 
-  unsigned optLevel = 0;
-  if (llvmO1)
-    optLevel = 1;
-  if (llvmO2)
-    optLevel = 2;
-  if (llvmO3)
-    optLevel = 3;
-
-  std::function<llvm::Error(llvm::Module *)> transformer;
-  if (llvmPasses.getValue().empty())
-    transformer = mlir::makeOptimizingTransformer(optLevel, /*sizeLevel=*/0);
-  else
-    transformer = mlir::makeLLVMPassesTransformer(llvmPasses.getValue());
-
+  auto transformer = mlir::makeLLVMPassesTransformer(llvmPasses.getValue());
   auto error = compileAndExecute(m.get(), mainFuncName.getValue(), transformer);
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),