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:
PyExecutionEngine(MlirExecutionEngine executionEngine)22 PyExecutionEngine(MlirExecutionEngine executionEngine)
23 : executionEngine(executionEngine) {}
PyExecutionEngine(PyExecutionEngine && other)24 PyExecutionEngine(PyExecutionEngine &&other)
25 : executionEngine(other.executionEngine) {
26 other.executionEngine.ptr = nullptr;
27 }
~PyExecutionEngine()28 ~PyExecutionEngine() {
29 if (!mlirExecutionEngineIsNull(executionEngine))
30 mlirExecutionEngineDestroy(executionEngine);
31 }
get()32 MlirExecutionEngine get() { return executionEngine; }
33
release()34 void release() {
35 executionEngine.ptr = nullptr;
36 referencedObjects.clear();
37 }
getCapsule()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.
addReferencedObject(pybind11::object obj)45 void addReferencedObject(pybind11::object obj) {
46 referencedObjects.push_back(obj);
47 }
48
createFromCapsule(pybind11::object capsule)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.
PYBIND11_MODULE(_mlirExecutionEngine,m)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 = mlirExecutionEngineLookup(
104 executionEngine.get(),
105 mlirStringRefCreate(func.c_str(), func.size()));
106 return reinterpret_cast<uintptr_t>(res);
107 },
108 "Lookup function `func` in the ExecutionEngine.")
109 .def(
110 "raw_register_runtime",
111 [](PyExecutionEngine &executionEngine, const std::string &name,
112 py::object callbackObj) {
113 executionEngine.addReferencedObject(callbackObj);
114 uintptr_t rawSym =
115 py::cast<uintptr_t>(py::getattr(callbackObj, "value"));
116 mlirExecutionEngineRegisterSymbol(
117 executionEngine.get(),
118 mlirStringRefCreate(name.c_str(), name.size()),
119 reinterpret_cast<void *>(rawSym));
120 },
121 py::arg("name"), py::arg("callback"),
122 "Register `callback` as the runtime symbol `name`.")
123 .def(
124 "dump_to_object_file",
125 [](PyExecutionEngine &executionEngine, const std::string &fileName) {
126 mlirExecutionEngineDumpToObjectFile(
127 executionEngine.get(),
128 mlirStringRefCreate(fileName.c_str(), fileName.size()));
129 },
130 "Dump ExecutionEngine to an object file.");
131 }
132