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() { 35 executionEngine.ptr = nullptr; 36 referencedObjects.clear(); 37 } 38 pybind11::object getCapsule() { 39 return py::reinterpret_steal<py::object>( 40 mlirPythonExecutionEngineToCapsule(get())); 41 } 42 43 // Add an object to the list of referenced objects whose lifetime must exceed 44 // those of the ExecutionEngine. 45 void addReferencedObject(pybind11::object obj) { 46 referencedObjects.push_back(obj); 47 } 48 49 static pybind11::object createFromCapsule(pybind11::object capsule) { 50 MlirExecutionEngine rawPm = 51 mlirPythonCapsuleToExecutionEngine(capsule.ptr()); 52 if (mlirExecutionEngineIsNull(rawPm)) 53 throw py::error_already_set(); 54 return py::cast(PyExecutionEngine(rawPm), py::return_value_policy::move); 55 } 56 57 private: 58 MlirExecutionEngine executionEngine; 59 // We support Python ctypes closures as callbacks. Keep a list of the objects 60 // so that they don't get garbage collected. (The ExecutionEngine itself 61 // just holds raw pointers with no lifetime semantics). 62 std::vector<py::object> referencedObjects; 63 }; 64 65 } // anonymous namespace 66 67 /// Create the `mlir.execution_engine` module here. 68 PYBIND11_MODULE(_mlirExecutionEngine, m) { 69 m.doc() = "MLIR Execution Engine"; 70 71 //---------------------------------------------------------------------------- 72 // Mapping of the top-level PassManager 73 //---------------------------------------------------------------------------- 74 py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local()) 75 .def(py::init<>([](MlirModule module, int optLevel, 76 const std::vector<std::string> &sharedLibPaths) { 77 llvm::SmallVector<MlirStringRef, 4> libPaths; 78 for (const std::string &path : sharedLibPaths) 79 libPaths.push_back({path.c_str(), path.length()}); 80 MlirExecutionEngine executionEngine = mlirExecutionEngineCreate( 81 module, optLevel, libPaths.size(), libPaths.data()); 82 if (mlirExecutionEngineIsNull(executionEngine)) 83 throw std::runtime_error( 84 "Failure while creating the ExecutionEngine."); 85 return new PyExecutionEngine(executionEngine); 86 }), 87 py::arg("module"), py::arg("opt_level") = 2, 88 py::arg("shared_libs") = py::list(), 89 "Create a new ExecutionEngine instance for the given Module. The " 90 "module must contain only dialects that can be translated to LLVM. " 91 "Perform transformations and code generation at the optimization " 92 "level `opt_level` if specified, or otherwise at the default " 93 "level of two (-O2). Load a list of libraries specified in " 94 "`shared_libs`.") 95 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, 96 &PyExecutionEngine::getCapsule) 97 .def("_testing_release", &PyExecutionEngine::release, 98 "Releases (leaks) the backing ExecutionEngine (for testing purpose)") 99 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule) 100 .def( 101 "raw_lookup", 102 [](PyExecutionEngine &executionEngine, const std::string &func) { 103 auto *res = mlirExecutionEngineLookupPacked( 104 executionEngine.get(), 105 mlirStringRefCreate(func.c_str(), func.size())); 106 return reinterpret_cast<uintptr_t>(res); 107 }, 108 py::arg("func_name"), 109 "Lookup function `func` in the ExecutionEngine.") 110 .def( 111 "raw_register_runtime", 112 [](PyExecutionEngine &executionEngine, const std::string &name, 113 py::object callbackObj) { 114 executionEngine.addReferencedObject(callbackObj); 115 uintptr_t rawSym = 116 py::cast<uintptr_t>(py::getattr(callbackObj, "value")); 117 mlirExecutionEngineRegisterSymbol( 118 executionEngine.get(), 119 mlirStringRefCreate(name.c_str(), name.size()), 120 reinterpret_cast<void *>(rawSym)); 121 }, 122 py::arg("name"), py::arg("callback"), 123 "Register `callback` as the runtime symbol `name`.") 124 .def( 125 "dump_to_object_file", 126 [](PyExecutionEngine &executionEngine, const std::string &fileName) { 127 mlirExecutionEngineDumpToObjectFile( 128 executionEngine.get(), 129 mlirStringRefCreate(fileName.c_str(), fileName.size())); 130 }, 131 py::arg("file_name"), "Dump ExecutionEngine to an object file."); 132 } 133