1 //===- ExecutionEngineModule.cpp - Python module for execution engine -----===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir-c/Bindings/Python/Interop.h" 10 #include "mlir-c/ExecutionEngine.h" 11 #include "mlir/Bindings/Python/PybindAdaptors.h" 12 13 namespace py = pybind11; 14 using namespace mlir; 15 using namespace mlir::python; 16 17 namespace { 18 19 /// Owning Wrapper around an ExecutionEngine. 20 class PyExecutionEngine { 21 public: 22 PyExecutionEngine(MlirExecutionEngine executionEngine) 23 : executionEngine(executionEngine) {} 24 PyExecutionEngine(PyExecutionEngine &&other) 25 : executionEngine(other.executionEngine) { 26 other.executionEngine.ptr = nullptr; 27 } 28 ~PyExecutionEngine() { 29 if (!mlirExecutionEngineIsNull(executionEngine)) 30 mlirExecutionEngineDestroy(executionEngine); 31 } 32 MlirExecutionEngine get() { return executionEngine; } 33 34 void release() { executionEngine.ptr = nullptr; } 35 pybind11::object getCapsule() { 36 return py::reinterpret_steal<py::object>( 37 mlirPythonExecutionEngineToCapsule(get())); 38 } 39 40 static pybind11::object createFromCapsule(pybind11::object capsule) { 41 MlirExecutionEngine rawPm = 42 mlirPythonCapsuleToExecutionEngine(capsule.ptr()); 43 if (mlirExecutionEngineIsNull(rawPm)) 44 throw py::error_already_set(); 45 return py::cast(PyExecutionEngine(rawPm), py::return_value_policy::move); 46 } 47 48 private: 49 MlirExecutionEngine executionEngine; 50 }; 51 52 } // anonymous namespace 53 54 /// Create the `mlir.execution_engine` module here. 55 PYBIND11_MODULE(_mlirExecutionEngine, m) { 56 m.doc() = "MLIR Execution Engine"; 57 58 //---------------------------------------------------------------------------- 59 // Mapping of the top-level PassManager 60 //---------------------------------------------------------------------------- 61 py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local()) 62 .def(py::init<>([](MlirModule module, int optLevel, 63 const std::vector<std::string> &sharedLibPaths) { 64 llvm::SmallVector<MlirStringRef, 4> libPaths; 65 for (const std::string &path : sharedLibPaths) 66 libPaths.push_back({path.c_str(), path.length()}); 67 MlirExecutionEngine executionEngine = mlirExecutionEngineCreate( 68 module, optLevel, libPaths.size(), libPaths.data()); 69 if (mlirExecutionEngineIsNull(executionEngine)) 70 throw std::runtime_error( 71 "Failure while creating the ExecutionEngine."); 72 return new PyExecutionEngine(executionEngine); 73 }), 74 py::arg("module"), py::arg("opt_level") = 2, 75 py::arg("shared_libs") = py::list(), 76 "Create a new ExecutionEngine instance for the given Module. The " 77 "module must contain only dialects that can be translated to LLVM. " 78 "Perform transformations and code generation at the optimization " 79 "level `opt_level` if specified, or otherwise at the default " 80 "level of two (-O2). Load a list of libraries specified in " 81 "`shared_libs`.") 82 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, 83 &PyExecutionEngine::getCapsule) 84 .def("_testing_release", &PyExecutionEngine::release, 85 "Releases (leaks) the backing ExecutionEngine (for testing purpose)") 86 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule) 87 .def( 88 "raw_lookup", 89 [](PyExecutionEngine &executionEngine, const std::string &func) { 90 auto *res = mlirExecutionEngineLookup( 91 executionEngine.get(), 92 mlirStringRefCreate(func.c_str(), func.size())); 93 return reinterpret_cast<uintptr_t>(res); 94 }, 95 "Lookup function `func` in the ExecutionEngine.") 96 .def( 97 "raw_register_runtime", 98 [](PyExecutionEngine &executionEngine, const std::string &name, 99 uintptr_t sym) { 100 mlirExecutionEngineRegisterSymbol( 101 executionEngine.get(), 102 mlirStringRefCreate(name.c_str(), name.size()), 103 reinterpret_cast<void *>(sym)); 104 }, 105 "Lookup function `func` in the ExecutionEngine.") 106 .def( 107 "dump_to_object_file", 108 [](PyExecutionEngine &executionEngine, const std::string &fileName) { 109 mlirExecutionEngineDumpToObjectFile( 110 executionEngine.get(), 111 mlirStringRefCreate(fileName.c_str(), fileName.size())); 112 }, 113 "Dump ExecutionEngine to an object file."); 114 } 115