[MLIR][PYTHON] Provide opt level for ExecutionEngine Python binding

[bondhugula]

Summary

Provide an option to specify optimization level when creating an
ExecutionEngine via the MLIR JIT Python binding. Not only is the
specified optimization level used for code generation, but all LLVM
optimization passes at the optimization level are also run prior to
machine code generation (akin to the mlir-cpu-runner tool).

Default opt level continues to remain at level two (-O2).

Contributions in part from Prashant Kumar <[e-mail]> as well.

[bondhugula]

Update commit summary.

[stellaraccident]

Thanks for the patch! I'm accepting this to help with timezones but would like my comments addressed/further discussed before landing.

mlir/include/mlir-c/ExecutionEngine.h

44

I think as a general rule, we're not trying to carry forward unsigned into new APIs. For a general quantity (not on the order of pointer/system memory sized), prefer int. Here and elsewhere in this patch.

mlir/lib/Bindings/Python/ExecutionEngine.cpp

62

You can write this as:

.def(py::init<>([](PyModule &module, int optLevel) {
  ...
}, py::arg("module"), py::arg("opt_level") = 2,
"My docstring")

For trivial cases of defaulting arguments like this, it is much better to have a python level default arg vs an overload.

85

When you remove this overload and adapt the docstring, note that any python visible args should be in snake_case.

[bondhugula]

Address comments from @stellaraccident.

mlir/lib/Bindings/Python/ExecutionEngine.cpp

62

Thanks very much! Updated.

[bondhugula]

unsigned -> int for optLevel.

[bondhugula]

@stellaraccident Leaving this for your second look.

[stellaraccident]

Thanks!

mlir/include/mlir-c/ExecutionEngine.h

 DEFINE_C_API_STRUCT(MlirExecutionEngine, void);
 
 #undef DEFINE_C_API_STRUCT
 
 /// Creates an ExecutionEngine for the provided ModuleOp. The ModuleOp is
 /// expected to be "translatable" to LLVM IR (only contains operations in
 /// dialects that implement the `LLVMTranslationDialectInterface`). The module
 /// ownership stays with the client and can be destroyed as soon as the call
-/// returns.
-/// TODO: figure out options (optimization level, etc.).
-MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op);
+/// returns. `optLevel` is the optimization level to be used for transformation
+/// and code generation. LLVM passes at `optLevel` are run before code
+/// generation.
+/// TODO: figure out other options.
+MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op,
+                                                                 int optLevel);
 
 /// Destroy an ExecutionEngine instance.
 MLIR_CAPI_EXPORTED void mlirExecutionEngineDestroy(MlirExecutionEngine jit);
 
 /// Checks whether an execution engine is null.
 static inline bool mlirExecutionEngineIsNull(MlirExecutionEngine jit) {
   return !jit.ptr;
 }

mlir/lib/Bindings/Python/ExecutionEngine.cpp

 } // anonymous namespace
 
 /// Create the `mlir.execution_engine` module here.
 void mlir::python::populateExecutionEngineSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
   py::class_<PyExecutionEngine>(m, "ExecutionEngine")
-      .def(py::init<>([](PyModule &module) {
+      .def(py::init<>([](PyModule &module, int optLevel) {
              MlirExecutionEngine executionEngine =
-                 mlirExecutionEngineCreate(module.get());
+                 mlirExecutionEngineCreate(module.get(), optLevel);
              if (mlirExecutionEngineIsNull(executionEngine))
                throw std::runtime_error(
                    "Failure while creating the ExecutionEngine.");
              return new PyExecutionEngine(executionEngine);
            }),
+           py::arg("module"), py::arg("opt_level") = 2,
            "Create a new ExecutionEngine instance for the given Module. The "
-           "module must "
-           "contain only dialects that can be translated to LLVM.")
+           "module must contain only dialects that can be translated to LLVM. "
+           "Perform transformations and code generation at the optimization "
+           "level `opt_level` if specified, or otherwise at the default "
+           "level of two (-O2).")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyExecutionEngine::getCapsule)
       .def("_testing_release", &PyExecutionEngine::release,
            "Releases (leaks) the backing ExecutionEngine (for testing purpose)")
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule)
       .def(
           "raw_lookup",
           [](PyExecutionEngine &executionEngine, const std::string &func) {
             auto *res = mlirExecutionEngineLookup(
                 executionEngine.get(),
                 mlirStringRefCreate(func.c_str(), func.size()));
             return reinterpret_cast<uintptr_t>(res);
           },
           "Lookup function `func` in the ExecutionEngine.")
       .def(
           "raw_register_runtime",
           [](PyExecutionEngine &executionEngine, const std::string &name,
              uintptr_t sym) {

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

 //===- ExecutionEngine.cpp - C API for MLIR JIT ---------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/ExecutionEngine.h"
 #include "mlir/CAPI/ExecutionEngine.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
+#include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "llvm/ExecutionEngine/Orc/Mangling.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
-extern "C" MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op) {
-  static bool init_once = [] {
+extern "C" MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op,
+                                                         int optLevel) {
+  static bool initOnce = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
     return true;
   }();
-  (void)init_once;
+  (void)initOnce;
 
   mlir::registerLLVMDialectTranslation(*unwrap(op)->getContext());
-  auto jitOrError = ExecutionEngine::create(unwrap(op));
+
+  auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+  if (!tmBuilderOrError) {
+    llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
+    return MlirExecutionEngine{nullptr};
+  }
+  auto tmOrError = tmBuilderOrError->createTargetMachine();
+  if (!tmOrError) {
+    llvm::errs() << "Failed to create a TargetMachine for the host\n";
+    return MlirExecutionEngine{nullptr};
+  }
+
+  // Create a transformer to run all LLVM optimization passes at the
+  // specified optimization level.
+  auto llvmOptLevel = static_cast<llvm::CodeGenOpt::Level>(optLevel);
+  auto transformer = mlir::makeLLVMPassesTransformer(
+      /*passes=*/{}, llvmOptLevel, /*targetMachine=*/tmOrError->get());
+  auto jitOrError = ExecutionEngine::create(
+      unwrap(op), /*llvmModuleBuilder=*/{}, transformer, llvmOptLevel);
   if (!jitOrError) {
     consumeError(jitOrError.takeError());
     return MlirExecutionEngine{nullptr};
   }
   return wrap(jitOrError->release());
 }
 
 extern "C" void mlirExecutionEngineDestroy(MlirExecutionEngine jit) {

mlir/test/CAPI/execution_engine.c

 "  func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {     \n"
 "    %res = std.addi %arg0, %arg0 : i32                                     \n"
 "    return %res : i32                                                      \n"
 "  }                                                                        \n"
 "}"));
   // clang-format on
   lowerModuleToLLVM(ctx, module);
   mlirRegisterAllLLVMTranslations(ctx);
-  MlirExecutionEngine jit = mlirExecutionEngineCreate(module);
+  MlirExecutionEngine jit = mlirExecutionEngineCreate(module, /*optLevel=*/2);
   if (mlirExecutionEngineIsNull(jit)) {
     fprintf(stderr, "Execution engine creation failed");
     exit(2);
   }
   int input = 42;
   int result = -1;
   void *args[2] = {&input, &result};
   if (mlirLogicalResultIsFailure(mlirExecutionEngineInvokePacked(