ExecutionEngine: provide utils for running CLI-configured LLVM passes

A recent change introduced a possibility to run LLVM IR transformation during
JIT-compilation in the ExecutionEngine.  Provide helper functions that
construct IR transformers given either clang-style optimization levels or a
list passes to run.  The latter wraps the LLVM command line option parser to
parse strings rather than actual command line arguments.  As a result, we can
run either of

    mlir-cpu-runner -O3 input.mlir
    mlir-cpu-runner -some-mlir-pass -llvm-opts="-llvm-pass -other-llvm-pass"

to combine different transformations.  The transformer builder functions are
provided as a separate library that depends on LLVM pass libraries unlike the
main execution engine library.  The library can be used for integrating MLIR
execution engine into external frameworks.

PiperOrigin-RevId: 234173493

mlir/include/mlir/ExecutionEngine/OptUtils.h

@@ -0,0 +1,55 @@
+//===- OptUtils.h - MLIR Execution Engine opt pass utilities ----*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file declares the utility functions to trigger LLVM optimizations from
+// MLIR Execution Engine.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_EXECUTIONENGINE_OPTUTILS_H_
+#define MLIR_EXECUTIONENGINE_OPTUTILS_H_
+
+#include <functional>
+#include <string>
+
+namespace llvm {
+class Module;
+class Error;
+} // namespace llvm
+
+namespace mlir {
+
+/// Initialize LLVM passes that can be when running MLIR code using
+/// ExecutionEngine.
+void initializeLLVMPasses();
+
+/// Create a module transformer function for MLIR ExecutionEngine that runs
+/// LLVM IR passes corresponding to the given speed and size optimization
+/// levels (e.g. -O2 or -Os).
+std::function<llvm::Error(llvm::Module *)>
+makeOptimizingTransformer(unsigned optLevel, unsigned sizeLevel);
+
+/// Create a module transformer function for MLIR ExecutionEngine that runs
+/// LLVM IR passes specified by the configuration string that uses the same
+/// syntax as LLVM opt tool.  For example, "-loop-distribute -loop-vectorize"
+/// will run the loop distribution pass followed by the loop vectorizer.
+std::function<llvm::Error(llvm::Module *)>
+makeLLVMPassesTransformer(std::string config);
+
+} // end namespace mlir
+
+#endif // LIR_EXECUTIONENGINE_OPTUTILS_H_

mlir/lib/ExecutionEngine/OptUtils.cpp

@@ -0,0 +1,118 @@
+//===- OptUtils.cpp - MLIR Execution Engine optimization pass utilities ---===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements the utility functions to trigger LLVM optimizations from
+// MLIR Execution Engine.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/ExecutionEngine/OptUtils.h"
+
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/LegacyPassNameParser.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/StringSaver.h"
+#include "llvm/Transforms/IPO.h"
+#include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include <mutex>
+
+// 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) {
+  for (auto &func : m) {
+    funcPM.run(func);
+  }
+  modulePM.run(m);
+}
+
+// Initialize basic LLVM transformation passes under lock.
+void mlir::initializeLLVMPasses() {
+  static std::mutex mutex;
+  std::lock_guard<std::mutex> lock(mutex);
+
+  auto &registry = *llvm::PassRegistry::getPassRegistry();
+  llvm::initializeCore(registry);
+  llvm::initializeTransformUtils(registry);
+  llvm::initializeScalarOpts(registry);
+  llvm::initializeIPO(registry);
+  llvm::initializeInstCombine(registry);
+  llvm::initializeAggressiveInstCombine(registry);
+  llvm::initializeAnalysis(registry);
+  llvm::initializeVectorization(registry);
+}
+
+// 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);
+    runPasses(modulePM, funcPM, *m);
+
+    return llvm::Error::success();
+  };
+}
+
+// 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"));
+    llvm::BumpPtrAllocator allocator;
+    llvm::StringSaver saver(allocator);
+    llvm::SmallVector<const char *, 16> args;
+    args.push_back(""); // inject dummy program name
+    llvm::cl::TokenizeGNUCommandLine(config, saver, args);
+    llvm::cl::ParseCommandLineOptions(args.size(), args.data());
+
+    llvm::legacy::PassManager modulePM;
+
+    for (const auto *passInfo : llvmPasses) {
+      if (!passInfo->getNormalCtor())
+        continue;
+
+      auto *pass = passInfo->createPass();
+      if (!pass)
+        return llvm::make_error<llvm::StringError>(
+            "could not create pass " + passInfo->getPassName(),
+            llvm::inconvertibleErrorCode());
+
+      modulePM.add(pass);
+    }
+
+    modulePM.run(*m);
+    return llvm::Error::success();
+  };
+}

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

@@ -1,5 +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
 
 func @fabsf(f32) -> f32
 
@@ -27,4 +29,4 @@ func @foo(%a : memref<1x1xf32>) -> memref<1x1xf32> {
   return %a : memref<1x1xf32>
 }
 // NOMAIN: 2.234000e+03
-// NOMAIN-NEXT: 2.234000e+03
+// NOMAIN-NEXT: 2.234000e+03

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

@@ -22,6 +22,7 @@
 
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/ExecutionEngine/MemRefUtils.h"
+#include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
@@ -53,6 +54,23 @@ static llvm::cl::opt<std::string>
                  llvm::cl::value_desc("<function name>"),
                  llvm::cl::init("main"));
 
+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.
@@ -108,7 +126,9 @@ static void printMemRefArguments(ArrayRef<Type> argTypes,
   }
 }
 
-static Error compileAndExecute(Module *module, StringRef entryPoint) {
+static Error
+compileAndExecute(Module *module, StringRef entryPoint,
+                  std::function<llvm::Error(llvm::Module *)> transformer) {
   Function *mainFunction = module->getNamedFunction(entryPoint);
   if (!mainFunction || mainFunction->getBlocks().empty()) {
     return make_string_error("entry point not found");
@@ -128,7 +148,7 @@ static Error compileAndExecute(Module *module, StringRef entryPoint) {
   if (!expectedArguments)
     return expectedArguments.takeError();
 
-  auto expectedEngine = mlir::ExecutionEngine::create(module);
+  auto expectedEngine = mlir::ExecutionEngine::create(module, transformer);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
@@ -150,7 +170,14 @@ 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();
 
   MLIRContext context;
   auto m = parseMLIRInput(inputFilename, &context);
@@ -158,7 +185,22 @@ int main(int argc, char **argv) {
     llvm::errs() << "could not parse the input IR\n";
     return 1;
   }
-  auto error = compileAndExecute(m.get(), mainFuncName.getValue());
+
+  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 error = compileAndExecute(m.get(), mainFuncName.getValue(), transformer);
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),
                         [&exitCode](const llvm::ErrorInfoBase &info) {