Set mlir-cpu-runner JIT codegen opt level correctly

- the JIT codegen was being run at the default -O0 level; instead,
  propagate the opt level from the cmd line.

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#123

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/123 from bondhugula:jit-runner 3b055e47f94c9a48bf487f6400787478738cda02
PiperOrigin-RevId: 267778586

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

@@ -72,13 +72,15 @@ public:
 
   /// Creates an execution engine for the given module.  If `transformer` is
   /// provided, it will be called on the LLVM module during JIT-compilation and
-  /// can be used, e.g., for reporting or optimization.
-  /// If `sharedLibPaths` are provided, the underlying JIT-compilation will open
-  /// and link the shared libraries for symbol resolution.
-  /// If `objectCache` is provided, JIT compiler will use it to store the object
-  /// generated for the given module.
+  /// can be used, e.g., for reporting or optimization. `jitCodeGenOptLevel`,
+  /// when provided, is used as the optimization level for target code
+  /// generation. If `sharedLibPaths` are provided, the underlying
+  /// JIT-compilation will open and link the shared libraries for symbol
+  /// resolution. If `objectCache` is provided, JIT compiler will use it to
+  /// store the object generated for the given module.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>> create(
       ModuleOp m, std::function<llvm::Error(llvm::Module *)> transformer = {},
+      Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
       ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = false);
 
   /// Looks up a packed-argument function with the given name and returns a

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

@@ -198,6 +198,7 @@ ExecutionEngine::ExecutionEngine(bool enableObjectCache)
 
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
     ModuleOp m, std::function<Error(llvm::Module *)> transformer,
+    Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
     ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache) {
   auto engine = std::make_unique<ExecutionEngine>(enableObjectCache);
 
@@ -264,6 +265,8 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   // LLJITWithObjectCache example.
   auto compileFunctionCreator = [&](JITTargetMachineBuilder JTMB)
       -> Expected<IRCompileLayer::CompileFunction> {
+    if (jitCodeGenOptLevel)
+      JTMB.setCodeGenOptLevel(jitCodeGenOptLevel.getValue());
     auto TM = JTMB.createTargetMachine();
     if (!TM)
       return TM.takeError();

mlir/lib/Support/JitRunner.cpp

@@ -81,14 +81,18 @@ static llvm::cl::list<const llvm::PassInfo *, bool, llvm::PassNameParser>
                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));
+static llvm::cl::opt<bool>
+    optO0("O0", llvm::cl::desc("Run opt passes and codegen at O0"),
+          llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool>
+    optO1("O1", llvm::cl::desc("Run opt passes and codegen at O1"),
+          llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool>
+    optO2("O2", llvm::cl::desc("Run opt passes and codegen at O2"),
+          llvm::cl::cat(optFlags));
+static llvm::cl::opt<bool>
+    optO3("O3", llvm::cl::desc("Run opt passes and codegen at O3"),
+          llvm::cl::cat(optFlags));
 
 static llvm::cl::OptionCategory clOptionsCategory("linking options");
 static llvm::cl::list<std::string>
@@ -178,14 +182,34 @@ static LogicalResult convertAffineStandardToLLVMIR(ModuleOp module) {
   return manager.run(module);
 }
 
+static llvm::Optional<unsigned> getCommandLineOptLevel() {
+  llvm::Optional<unsigned> optLevel;
+  llvm::SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
+      optO0, optO1, optO2, optO3};
+
+  // Determine if there is an optimization flag present.
+  for (unsigned j = 0; j < 4; ++j) {
+    auto &flag = optFlags[j].get();
+    if (flag) {
+      optLevel = j;
+      break;
+    }
+  }
+  return optLevel;
+}
+
 // JIT-compile the given module and run "entryPoint" with "args" as arguments.
 static Error
 compileAndExecute(ModuleOp module, StringRef entryPoint,
                   std::function<llvm::Error(llvm::Module *)> transformer,
                   void **args) {
+  Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
+  if (auto clOptLevel = getCommandLineOptLevel())
+    jitCodeGenOptLevel =
+        static_cast<llvm::CodeGenOpt::Level>(clOptLevel.getValue());
   SmallVector<StringRef, 4> libs(clSharedLibs.begin(), clSharedLibs.end());
-  auto expectedEngine =
-      mlir::ExecutionEngine::create(module, transformer, libs);
+  auto expectedEngine = mlir::ExecutionEngine::create(module, transformer,
+                                                      jitCodeGenOptLevel, libs);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
@@ -296,26 +320,24 @@ int mlir::JitRunnerMain(
   initializeLLVM();
   mlir::initializeLLVMPasses();
 
-  llvm::SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
-      optO0, optO1, optO2, optO3};
-
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
-  llvm::SmallVector<const llvm::PassInfo *, 4> passes;
-  llvm::Optional<unsigned> optLevel;
+  llvm::Optional<unsigned> optLevel = getCommandLineOptLevel();
+  llvm::SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
+      optO0, optO1, optO2, optO3};
   unsigned optCLIPosition = 0;
   // Determine if there is an optimization flag present, and its CLI position
   // (optCLIPosition).
   for (unsigned j = 0; j < 4; ++j) {
     auto &flag = optFlags[j].get();
     if (flag) {
-      optLevel = j;
       optCLIPosition = flag.getPosition();
       break;
     }
   }
   // Generate vector of pass information, plus the index at which we should
   // insert any optimization passes in that vector (optPosition).
+  llvm::SmallVector<const llvm::PassInfo *, 4> passes;
   unsigned optPosition = 0;
   for (unsigned i = 0, e = llvmPasses.size(); i < e; ++i) {
     passes.push_back(llvmPasses[i]);