diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h --- a/mlir/include/mlir-c/ExecutionEngine.h +++ b/mlir/include/mlir-c/ExecutionEngine.h @@ -36,9 +36,12 @@ /// 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 --- a/mlir/lib/Bindings/Python/ExecutionEngine.cpp +++ b/mlir/lib/Bindings/Python/ExecutionEngine.cpp @@ -59,17 +59,20 @@ // Mapping of the top-level PassManager //---------------------------------------------------------------------------- py::class_(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 --- 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(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 --- a/mlir/test/CAPI/execution_engine.c +++ b/mlir/test/CAPI/execution_engine.c @@ -48,7 +48,7 @@ // 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);