[Mlir-commits] [mlir] 185ce8c - [MLIR][PYTHON] Provide opt level for ExecutionEngine Python binding

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-05-16T13:58:49+05:30
New Revision: 185ce8cdfc7dba6cb6611b04346bfa03d008d5bf

URL: https://github.com/llvm/llvm-project/commit/185ce8cdfc7dba6cb6611b04346bfa03d008d5bf
DIFF: https://github.com/llvm/llvm-project/commit/185ce8cdfc7dba6cb6611b04346bfa03d008d5bf.diff

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

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 <prashantk at polymagelabs.com>
as well.

Differential Revision: https://reviews.llvm.org/D102551

Added: 
    

Modified: 
    mlir/include/mlir-c/ExecutionEngine.h
    mlir/lib/Bindings/Python/ExecutionEngine.cpp
    mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
    mlir/test/CAPI/execution_engine.c

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h
index 4a5d6ad9f648f..289e8f73dd4ac 100644
--- a/mlir/include/mlir-c/ExecutionEngine.h
+++ b/mlir/include/mlir-c/ExecutionEngine.h
@@ -36,9 +36,12 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void);
 /// 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);

diff  --git a/mlir/lib/Bindings/Python/ExecutionEngine.cpp b/mlir/lib/Bindings/Python/ExecutionEngine.cpp
index b5c8dde7563f5..38cf6b2ca4e8c 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngine.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngine.cpp
@@ -59,17 +59,20 @@ 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,

diff  --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
index 36f24ed889650..dfde38aee9698 100644
--- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
@@ -10,22 +10,42 @@
 #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};

diff  --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c
index 160bdbebe9e53..a2af7e8591057 100644
--- a/mlir/test/CAPI/execution_engine.c
+++ b/mlir/test/CAPI/execution_engine.c
@@ -48,7 +48,7 @@ void testSimpleExecution() {
   // 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);